Equation Sets

build_nn_function also accepts a whole set of equations: an EquationSet, an arbitrarily nested NamedTuple whose leaves are symbolic expressions. The result is one function whose output has the same nesting.

It also accepts a NetworkParameters of expressions, which is what a symbolic derivative is — a derivative with respect to a network's parameters has the shape of those parameters. The two are separate methods rather than one signature over a union, because they are separate things: an equation set is a bundle a caller wrote, and a parameter set is what a network is evaluated at.

using SymbolicNeuralNetworks
using AbstractNeuralNetworks: Chain, Dense, NeuralNetwork, params
import Random
Random.seed!(123)

c = Chain(Dense(2, 3, tanh), Dense(3, 2, tanh))
snn = SymbolicNeuralNetwork(c)
nn = NeuralNetwork(c)
ps = params(nn)

eqs = (output = c(snn.input, params(snn)),
       squared = c(snn.input, params(snn)) .^ 2,
       total = sum(c(snn.input, params(snn))))

f = build_nn_function(eqs, params(snn), snn.input)
f([1.0, 2.0], ps)

(output = [0.7253418818031555, 0.01738299849106655], squared = [0.5261208454977429, 0.0003021686365404219], total = 0.742724880294222)

This is not just a convenience. The entries are generated as a single function, so everything they have in common is computed once:

f(rand(2, 4), ps)

(output = [0.865800656715577 0.8698088387104005 0.8800500114317996 0.8818170586424531; 0.4714584498475237 0.5173332250649776 0.33500085876100455 0.2023958091157365], squared = [0.7496107771691245 0.7565674158987354 0.7744880226211106 0.7776013249128275; 0.22227306993263005 0.26763366575613073 0.11222557537061052 0.04096406354761365], total = [1.3372591065631008 1.3871420637753782 1.2150508701928042 1.0842128677581897])

The entries follow the shape rules of Building Functions individually — a vector-valued entry becomes a matrix over a batch, a scalar-valued one a $1\times{}N$ matrix — and each is copied out of the joint result, so the entries never alias one another.

Why a single function

The main use of this is a symbolic gradient, whose entries are the derivatives with respect to each parameter array and therefore all share the entire forward pass:

using SymbolicNeuralNetworks: symbolic_parameter_gradient
using LinearAlgebra: norm

gradient = symbolic_parameter_gradient(norm(c(snn.input, params(snn))) ^ 2, snn)
keys(gradient), keys(gradient.L1)

((:L1, :L2), (:W, :b))

build_nn_function(gradient, params(snn), snn.input; reduce = +)(rand(2, 8), ps).L2.W

2×3 Matrix{Float64}:
 -2.23241  0.736642  -2.41855
 -3.36492  1.29202   -3.62235

Building each of those four entries as its own function would re-derive the forward pass four times and compile four RuntimeGeneratedFunctions instead of one. Internally the set is flattened into a single vector of scalar equations by flatten_equations and the flat result is split up again by split_result.

Arrays of equation sets

An array of equation sets — what Gradient produces for an array-valued expression — gives a function returning an array of results. Each entry of the array is built jointly; the entries themselves are independent and stay separate functions.

using SymbolicNeuralNetworks: Gradient, derivative

g = build_nn_function(derivative(Gradient(snn)), params(snn), snn.input)
result = g([1.0, 2.0], ps)
length(result), result[1].L1.b

(2, [-0.06800030240708374, -0.0020276516911752938, -0.2956673045107465])