API

BATCH_INDEX

The name the generated kernels give the index of the batch column they evaluate. See index_by_batch.

FIXED_NAMES

The three names a generated kernel always gives its own arguments, whatever its arity. See is_reserved_name for the rest.

OUTPUT_NAME

The name the generated in-place kernels give the array they write into. See accumulate_into_output.

PARAMETER_NAME

The name the generated kernels give their parameter argument. The rewrite rules turn every parameter argument of the code Symbolics.build_function emits into a getproperty on it, so that the kernel can be called with a single NetworkParameters.

SymbolicExpression

Everything this package accepts as an equation: a single symbolic expression or an array of them.

Equations are built by applying a model to symbolic variables, e.g.

c(nn.input, params(nn))

Both representations Symbolics offers are covered. This package itself only ever produces arrays of scalar expressions (see symbolic_variables), but a Symbolics.Arr handed in by a user is accepted as well and normalised by scalar_expressions.

SymbolicVariables

A symbolic array that a generated function takes as a data argument: the network input, or the target output of a loss.

Both the Vector{Num} that SymbolicNeuralNetwork builds and the Symbolics.Arr that @variables x[1:n] produces are accepted; the latter is scalarised by scalar_expressions.

(::FeedForwardLoss)(model, params, input, output)

The FeedForwardLoss of AbstractNeuralNetworks, evaluated on symbolic arguments.

AbstractBatchedFunction{NDATA, R}

An executable function built from a symbolic equation, as returned by build_nn_function.

It wraps a kernel that evaluates a single sample (see build_kernel and build_kernel!) and adds everything the kernel does not do: iterating over a batch, allocating and shaping the result, and accepting a single sample or a three-dimensional batch instead of a matrix.

NDATA is the number of data arguments, so an instance is called as

f(input, ps)            # NDATA = 1
f(input, output, ps)    # NDATA = 2
f(x1, …, xNDATA, ps)    # in general

and R is how the per-sample results are combined — hcat or +.

There is no bound on NDATA. One data argument is the network input, a second is typically the target output of a loss, and the layerwise pullback uses one per entry of a layer's seam plus one for the output sensitivities — see seam_interface.

Result shapes

For an equation of size $(m, n, \ldots)$ and a batch of $N$ samples:

A scalar-valued equation counts as $m = 1$ here, so batching it with hcat gives a $1\times{}N$ matrix.

All data arguments must have the same number of dimensions and the same batch size.

Derivative

Supertype of the symbolic derivatives this package computes: Jacobian (with respect to the input of a network) and Gradient (with respect to its parameters). Use derivative to get the symbolic expression out of one.

EquationSet

A keyed set of symbolic expressions: a NamedTuple whose values are expressions, arrays of them, or further sets. It is what a caller writes by hand and hands to build_nn_function —

eqs = (a = c(nn.input, params(nn)), b = c(nn.input, params(nn)) .^ 2)

— and it is not a set of parameters, which is a NeuralNetworkParameters.NetworkParameters. A plain link and not an @extref: this package configures no DocumenterInterLinks inventories, so an @extref here has nothing to resolve against and fails the build at CrossReferences. The two share a shape and nothing else: a parameter set is the thing a network is evaluated at, and an equation set is a bundle of expressions that happen to be keyed. Naming them apart is what keeps a signature honest about which it wants; the one place both arrive is a symbolic gradient, which is parameter-shaped and holds expressions, and flatten_equations has a method for each.

An alias for NamedTuple rather than a type of its own, so a method taking one is a method on Base.NamedTuple — permissible only because the functions here are this package's. Do not extend a foreign generic on it.

EquationSetArrayFunction{NDATA}(functions)

The function build_nn_function returns for an array of EquationSets: it evaluates one EquationSetFunction per entry and collects the results.

EquationSetFunction{NDATA}(f, layout)

The function build_nn_function returns for an EquationSet: it evaluates the jointly generated f and puts the flat result back into the nesting recorded in layout.

FlatParameterFunction{NDATA}(f, layout)

The function build_flat_function returns: it lays its flat parameter argument out into the shape layout describes and calls f with it.

Gradient <: Derivative

Computes and stores the derivative of a symbolic expression with respect to the parameters of a SymbolicNeuralNetwork.

Constructors

Gradient(f, nn)

Differentiate the symbolic f with respect to the parameters of nn.

Gradient(nn)

Differentiate the symbolic output of nn, i.e. nn.model(nn.input, params(nn)).

Examples

using SymbolicNeuralNetworks: SymbolicNeuralNetwork, Gradient, derivative
using AbstractNeuralNetworks

c = Chain(Dense(2, 1, tanh))
nn = SymbolicNeuralNetwork(c)
(Gradient(nn) |> derivative)[1].L1.b

# output

1-element Vector{Symbolics.Num}:
 1 - (tanh(W_2₁ + W_1₁ˏ₁*x₁ + W_1₁ˏ₂*x₂)^2)

Implementation

Internally this uses symbolic_parameter_gradient. For an array-valued f the result is an array of the same shape whose entries are the parameter-shaped gradients of the corresponding entry of f — so the gradient of a matrix is a matrix of NetworkParameters, each of which is the ordinary gradient of one matrix element.

InPlaceBatchedFunction{NDATA}(kernel!, equation_size, reduction)

An AbstractBatchedFunction that allocates the result once and lets its kernel write every batch column into it. This is the default of build_nn_function; it costs a single allocation per call rather than one per column, at the price of not being differentiable by Zygote (Mutating arrays is not supported). Forward-mode AD works either way, as the element type of the preallocated array is promoted over the inputs — see promoted_eltype.

Jacobian <: Derivative

Computes and stores the derivative of a symbolic expression with respect to the input of a SymbolicNeuralNetwork.

Constructors

Jacobian(f, nn)
Jacobian(nn)

Differentiate the symbolic f with respect to the input of nn. If f is not supplied it is taken to be the symbolic output of the network, nn.model(nn.input, params(nn)).

Fields

1. f: the symbolic expression that was differentiated,
2. □: the symbolic Jacobian,
3. nn: the SymbolicNeuralNetwork.

Implementation

For a function $f:\mathbb{R}^n\to\mathbb{R}^m$ we use the convention

\[\square_{ij} = \frac{\partial}{\partial{}x_j}f_i, \text{ i.e. } \square \in \mathbb{R}^{m\times{}n},\]

which is also the one Zygote and ForwardDiff use. An f that is not a vector is flattened with vec first, so the rows of □ are indexed by vec(f); a scalar f gives a $1\times{}n$ Jacobian, i.e. its gradient with respect to the input as a row.

Examples

Here we compute the Jacobian of a single-layer neural network $x \mapsto \mathrm{tanh}(Wx + b)$, whose element-wise derivative is

\[ \frac{\partial}{\partial{}x_i}\sigma\left(\sum_{k}w_{jk}x_k + b_j\right) = \sigma'\left(\sum_{k}w_{jk}x_k + b_j\right)w_{ji},\]

and compare it to that expression. Note that $\mathrm{tanh}'(x) = \frac{4e^{2x}}{(e^{2x} + 1)^2}.$

using SymbolicNeuralNetworks
using SymbolicNeuralNetworks: Jacobian, derivative
using AbstractNeuralNetworks: Dense, Chain, NeuralNetwork, params
import Random

Random.seed!(123)

input_dim = 5
output_dim = 2
c = Chain(Dense(input_dim, output_dim, tanh))
nn = SymbolicNeuralNetwork(c)
jacobian = build_nn_function(derivative(Jacobian(nn)), nn)

ps = params(NeuralNetwork(c, Float64))
input = rand(input_dim)
Dtanh(x::Real) = 4 * exp(2 * x) / (1 + exp(2x)) ^ 2
analytic_jacobian(i, j) = Dtanh(sum(k -> ps.L1.W[j, k] * input[k], 1:input_dim) + ps.L1.b[j]) * ps.L1.W[j, i]
jacobian(input, ps) ≈ [analytic_jacobian(i, j) for j ∈ 1:output_dim, i ∈ 1:input_dim]

# output

true

LayerStep{Key}(layer, dλ, dθ)

One step of the adjoint sweep: the layer, the two generated functions the backward pass calls for it, and — as a type parameter — the key its parameters have in the parameter set.

Both functions take (x, λ, ps): the layer's own input, the sensitivity of the loss to the layer's output, and the parameters of the whole network. dλ returns the sensitivity with respect to the layer's input, dθ the derivative of the loss with respect to the layer's parameters.

A layer that carries data alongside the state takes that data in between, as further arguments before λ; seam_arguments is what produces them, and seam_interface the whole picture.

dλ is nothing for the first step of a sweep, which is the one place it is never called; see layer_step.

Taking the whole parameter set rather than the layer's own entry is what avoids a wrapper per call: the generated kernels were built from symbolic parameters nested under Key, so they reach for ps.<Key>.W and nothing has to be rebuilt for them.

Key is a type parameter rather than a field so that reading the layer's own entry out of the parameter set — step_parameters, which the forward pass needs — is an inferable getproperty on a name the compiler knows, rather than one on a Symbol it does not. Without it a chain whose layers hold differently shaped parameters infers their union.

LayerwiseGradientFunction{Keys}(steps, seed)

What a layerwise SymbolicPullback stores in place of a single generated gradient function.

Applying it to an input, a target output and the parameters runs the sweep and returns the derivative of the loss with respect to the parameters, as a NetworkParameters — the same thing the monolithic gradient function returns, so everything above it is unchanged.

Keys are the keys of the parameter set, as a type parameter so that the returned NamedTuple is inferable; see LayerStep for the same reasoning one level down.

OutOfPlaceBatchedFunction{NDATA}(kernel, equation_size, reduction)

An AbstractBatchedFunction that evaluates its kernel once per batch column and combines the results with Base.reduce. It allocates an array per column, but — unlike InPlaceBatchedFunction — it does not mutate anything and can therefore be differentiated by Zygote.

This is what build_nn_function returns for inplace = false, and for scalar-valued equations, for which Symbolics.build_function emits no in-place form.

ParameterGradient(gradient_function)

What a SymbolicPullback stores in its fun field. Applying it to an input, an output and the network parameters returns the PullbackFunction for those.

PassThroughLayer{N}()

A layer that returns its input unchanged, with no parameters. Used to obtain the loss as a function of the network's prediction rather than of its input; see loss_expression.

AbstractNeuralNetworks has no such interface: a NetworkLoss is applied as loss(model, ps, input, output), so the only way to ask it "what are you, as a function of the prediction and the target?" is to hand it a model whose prediction is its input.

PullbackFunction(gradient_function, input, output, parameters)

The function a SymbolicPullback returns as the second entry of its result. It takes the output sensitivities — which it ignores, as the loss is scalar-valued, see the extended help of SymbolicPullback — and returns the derivative of the loss with respect to the network parameters, as a NamedTuple.

SymbolicNeuralNetwork <: AbstractSymbolicNeuralNetwork

A symbolic representation of a (small) neural network.

It pairs a model with symbolic stand-ins for its parameters and its input, so that symbolic expressions can be built from it:

c = Chain(Dense(2, 1, tanh))
nn = SymbolicNeuralNetwork(c)
eq = c(nn.input, params(nn))

Those expressions can then be differentiated (Jacobian, Gradient, SymbolicPullback) and turned into executable code with build_nn_function.

Fields

• architecture: the neural network architecture,
• model: the model (typically a Chain that realizes the architecture),
• params: the symbolic parameters of the network, with the same nesting as the numeric ones,
• input: the symbolic input of the network, a Vector{Num}.

Constructors

SymbolicNeuralNetwork(nn)
SymbolicNeuralNetwork(arch, model)
SymbolicNeuralNetwork(model)
SymbolicNeuralNetwork(arch)

Build a SymbolicNeuralNetwork from an AbstractNeuralNetworks.NeuralNetwork, from an architecture and/or a model, or from a single layer.

Implementation

Parameters and input are built by symbolic_variables, i.e. they consist of scalar symbolic variables. The parameters are named W_1, W_2, … in the order in which they appear in the parameter set, the input entries x₁, x₂, ….

SymbolicPullback <: AbstractPullback

The symbolic pullback of a loss function: it evaluates the loss and the derivative of the loss with respect to the network parameters from generated code instead of by automatic differentiation.

Constructors

SymbolicPullback(nn, loss)
SymbolicPullback(nn)

Build the pullback of loss (an AbstractNeuralNetworks.NetworkLoss, by default a FeedForwardLoss) for the SymbolicNeuralNetwork nn.

Examples

using SymbolicNeuralNetworks
using AbstractNeuralNetworks
using AbstractNeuralNetworks: params
import Random
Random.seed!(123)

c = Chain(Dense(2, 1, tanh))
nn = NeuralNetwork(c)
snn = SymbolicNeuralNetwork(nn)
pb = SymbolicPullback(snn, FeedForwardLoss())
ps = params(nn)
typeof(pb(ps, nn.model, (rand(2), rand(1)))[2](1))

# output

@NamedTuple{L1::@NamedTuple{W::Matrix{Float64}, b::Vector{Float64}}}

Keyword Arguments

• layerwise: how the pullback is built (default :auto). With :auto it is composed layer by layer whenever that is the better choice (composes_layerwise), and built from one expression for the whole network otherwise. true demands the layerwise construction and errors if it does not apply, with a message naming the reason; false demands the monolithic one. See layerwise_gradient_function and monolithic_gradient_function — the two produce the same gradient, and differ by orders of magnitude in what it costs to build.
• cse: perform common subexpression elimination when generating code (default true). This matters most on the monolithic path: without it every one of the 2 * n_layers generated blocks re-emits the entire forward pass, which makes the code for networks with more than one hidden layer intractably large. See build_kernel.
• inplace: evaluate a batch with an in-place kernel (default true). The pullback is the end of the differentiation chain, so nothing differentiates through it and the default is what you want here; inplace = false exists for symmetry with build_nn_function.

Implementation

An instance stores

• loss: the NetworkLoss,
• fun: a ParameterGradient that produces the pullback function.

Calling the functor on ps, model and an (input, output) tuple returns

pullback.loss(model, ps, input, output), pullback.fun(input, output, ps)

where the second entry is again a function — of the output sensitivities.

Extended help

A NetworkLoss is scalar-valued, so the extra step of returning a function of the output sensitivities is not strictly necessary here — the equivalent of pb.fun(input, output, ps)(1) could be stored directly. It is however customary for a pullback to return a callable, which is why this package does so too.

The pullback is the derivative of a loss summed over the batch, which is why the generated function is built with reduce = +:

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

c = Chain(Dense(2, 1, tanh))
nn = NeuralNetwork(c)
snn = SymbolicNeuralNetwork(nn)
loss = FeedForwardLoss()
input_output = (rand(2), rand(1))

pb_values = SymbolicPullback(snn, loss)(params(nn), nn.model, input_output)[2](1)

soutput = Symbolics.variables(:y, 1:output_dimension(nn.model))
gradient = symbolic_parameter_gradient(loss(nn.model, params(snn), snn.input, soutput), snn)
pb_values2 = build_nn_function(gradient, params(snn), snn.input, soutput; reduce = +)(input_output..., params(nn))

pb_values == params(pb_values2)

# output

true

input_dimension(c::Chain)
output_dimension(c::Chain)

The dimensions a Chain maps between, taken from its first and last layer.

AbstractNeuralNetworks defines both for an AbstractLayer; these methods extend them to a whole Chain, which is what SymbolicNeuralNetwork needs to know how many symbolic input variables to build. They belong upstream too, see issue #35.

_assert_distinct_seam_variables(layer, sdata, sλ)

Reject a seam whose variables are not pairwise distinct, which is what a carried_variables that reuses the names layer_seed gives the state (x) or the sensitivities (λ) produces.

Nothing further down would notice. Symbolics.build_function binds one symbolic array to two argument slots and the generated code reads both of them from the last one, so the kernels are built, they run, and they return a gradient that is simply wrong — the one failure this construction must not have, and the reason checked_guess and represents_loss exist one level up.

Implementation

Here rather than in layer_seed, because checked_layer_seed catches everything layer_seed throws and turns it into a decline. A layer whose carried_variables collides is a bug in that layer, not a chain to quietly fall back on, so the check belongs on the far side of that try — and this is where the variables are handed to build_nn_function, which is where the damage would be done.

_assert_no_reserved_names_in_body(body)

Reject a generated body that already contains one of the names the kernels give their own arguments.

A symbolic variable that is passed to Symbolics.build_function becomes an argument and is renamed to ˍ₋argN, but one that is not — a variable left free in the equation, i.e. neither a data variable nor a parameter — survives into the body under its own name. If that name happens to be k, the kernel's batch index binds it and the equation silently evaluates with the column number in place of the variable; if it is ps, it binds the parameter set. Neither is caught by _assert_no_name_clash, which only sees the argument names.

The check has to run before the arguments are substituted, since afterwards the reserved names are all over the body legitimately. At that point the only symbols in the tree are ˍ₋argN, var"##cse#N" and literals — functions are embedded as objects — so a reserved name can only have come from a free variable.

_float_if_integer(T)

float(T) for integer types, T for everything else. See allocate_batch_output. Bool is included, as float(Bool) is Float64.

_rewrite_body(expression, data_names, parameter_paths, output_name)

Apply the rewrite rules of src/codegen/expression_rewriting.jl to the body of a generated function, in the order they depend on each other: the arguments are renamed first so that the later rules can recognise the data arguments, the array constructor is fixed before the batch index is added so that the typeof(…) it contains is not mistaken for a use of a data argument.

access_expression(name, path)

The expression that reads path out of the object called name, e.g. ps.L1.W for (:L1, :W).

accumulate_into_output(expr, output_name, reduction, equation_length)

Rewrite the out[i] = … assignments of in-place generated code so that a single preallocated array can hold the result of a whole batch.

Which rewrite applies depends on how the per-sample results are combined:

• reduction = +: the writes become +=, so every sample accumulates into the same buffer (which allocate_batch_output zeroes).
• reduction = hcat: the writes are shifted by $(k - 1)\cdot\mathrm{equation\_length}$, which is the offset of block k in the column-major layout of the concatenated result.

The generated code addresses its output with a single linear index whatever the shape of the equation, which is what makes that offset arithmetic correct. equation_length assignments are expected, one per entry of the equation; anything else throws, because a replace that matches nothing would leave a kernel that still compiles and runs but writes every sample to the same place.

Examples

using SymbolicNeuralNetworks: accumulate_into_output

accumulate_into_output(:(out[1] = a), :out, hcat, 1)

# output

:(out[1 + (k - 1) * 1] = a)

using SymbolicNeuralNetworks: accumulate_into_output

accumulate_into_output(:(out[1] = a), :out, +, 1)

# output

:(out[1] += a)

adjoint_step(steps, x, ps, seed, output)

The sensitivity of the loss to x and the parameter gradients of steps, given that x is what the remaining steps are applied to. See sweep.

The recursion bottoms out at the network's output, which the seed takes together with the target. So the chain's last layer has to return the model's output and nothing beside it — a layer that carries data through cannot be the last one, which is also what the monolithic construction and the loss itself require.

allocate_batch_output(T, equation_size, batch_size, reduction)

Allocate the result of evaluating an equation of size equation_size over batch_size samples, with the shape documented for AbstractBatchedFunction.

T is widened with float when it is an integer type: it comes from promoted_eltype, i.e. from the inputs, and an equation over integer inputs generally does not evaluate to an integer.

A scalar-valued equation — equation_size == () — counts as one of size $m = 1$, so hcat gives a $1\times{}N$ matrix and + a number. Those two methods are only reached for an empty batch: the in-place path, which is what allocates a result up front, does not exist for a scalar equation.

allocate_single_output(T, equation_size, reduction)

Allocate the result of evaluating an equation of size equation_size for a single sample. Unlike allocate_batch_output this keeps the shape of the equation, as the result may itself be a matrix.

argument_substitutions(generated_names, data_names, parameter_paths; output_name)

Map the argument names of the generated function onto the names the kernel uses, by position.

Symbolics.build_function was handed the data variables first and then one argument per parameter array (see parameter_arguments), so with output_name set — the in-place form prepends an output argument — the correspondence is

(ˍ₋out, ˍ₋arg1, …, ˍ₋arg_ndata, ˍ₋arg_{ndata+1}, …)  ↦  (out, x1, …, x_ndata, ps.L1.W, …)

Every parameter argument becomes a chain of getproperty calls on the single parameter argument ps, which is what lets the kernel be called with NetworkParameters instead of a flat argument list.

Examples

using SymbolicNeuralNetworks: argument_substitutions

substitutions = argument_substitutions([:ˍ₋arg1, :ˍ₋arg2], (:x1,), (((:L1, :W)),); output_name = nothing)
(substitutions[:ˍ₋arg1], substitutions[:ˍ₋arg2])

# output

(:x1, :(ps.L1.W))

batched(data)

data with any batch dimensions past the first collapsed into it, which is the shape the sweep evaluates in. For an input that carries data alongside the state, only the state is laid out; the rest is seam_arguments' business.

Implementation

build_nn_function accepts a result of $m\times{}N_1\times{}N_2$ as a batch with two batch dimensions (see AbstractBatchedFunction), and the monolithic construction of SymbolicPullback therefore evaluates one. The layerwise sweep cannot pass such an array to a layer — the forward pass is the layer called, and a Dense multiplies a matrix by a matrix — so the batch is laid out flat first.

That loses nothing. The pullback of a batch is the sum of the per-sample gradients, so how the samples are arranged cannot change it; the two shapes are checked against each other in test/derivatives/layerwise_pullback.jl. Nothing has to be restored afterwards either, for the same reason: what comes back is shaped like the parameters, not like the batch.

build_flat_function(eq, nn)
build_flat_function(eq, nn, soutput)
build_flat_function(eq, sparams, svariables...)

Turn a symbolic equation into an executable function whose parameter argument is a flat vector:

built_function(input, w)
built_function(input, output, w)

w may be a plain AbstractVector — in which case it is read through the layout of the parameters the equation was built from — or a NeuralNetworkParameters.FlatParameters, which carries its own.

Takes the same keyword arguments as build_nn_function, and accepts everything it accepts, including a NetworkParameters.

The third form — the symbolic parameters and the symbolic data variables given directly rather than taken from a network — is the one for degrees of freedom that are not a network's parameters: nothing here reads a model. It is what flat_parameter_gradient's second form pairs with. Unlike that one it needs a NetworkParameters specifically, because build_nn_function dispatches on one.

Examples

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

c = Chain(Dense(2, 1, tanh))
nn = NeuralNetwork(c)
snn = SymbolicNeuralNetwork(nn)
f = build_flat_function(c(snn.input, params(snn)), snn)

w, _ = flatten(params(nn))
f([1.0, 2.0], w) ≈ c([1.0, 2.0], params(nn))

# output

true

Implementation

The flat vector is unflattened and the ordinary generated function called on the result, rather than the code generation being changed to index into the vector. The conversion is a copyto! per leaf and costs a fraction of a forward pass; generating a different kernel would save that and lose the ability to hand the same equation a structured parameter set.

The conversion is unflatten and not unflatten! — it allocates the parameter set rather than writing into the one the layout was built from, so w may have a different element type, which is what lets ForwardDiff differentiate with respect to the flat form. This is unrelated to the inplace keyword, which says how the generated kernel evaluates a batch and is passed through untouched.

build_kernel!(equation, sparams, svariables...; reduction, cse)

Build an in-place kernel that writes the result for batch column k into a preallocated array:

kernel!(out, x1, ps, k)          # one data argument
kernel!(out, x1, x2, ps, k)      # two data arguments

Where in out the result goes depends on reduction; see accumulate_into_output.

Returns nothing for a scalar-valued equation, for which Symbolics.build_function emits no in-place form.

Evaluating a batch with such a kernel costs a single allocation instead of one array per column plus a Base.reduce fold, but the result is produced by mutation and can therefore not be differentiated by Zygote. See build_kernel for the keyword arguments.

build_kernel(equation, sparams, svariables...; cse)

Build an out-of-place kernel that evaluates equation for batch column k:

kernel(x1, ps, k)          # one data argument
kernel(x1, x2, ps, k)      # two data arguments

There is no bound on the number of data arguments; see data_name.

sparams are the symbolic parameters and svariables the symbolic data variables the equation was built from. See build_kernel! for the in-place counterpart and build_nn_function for the function that batching, allocation and reshaping are added to.

Examples

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

c = Chain(Dense(2, 1, tanh))
nn = NeuralNetwork(c)
snn = SymbolicNeuralNetwork(nn)
kernel = build_kernel(c(snn.input, params(snn)), params(snn), snn.input)
kernel([1.0 2.0; 3.0 4.0], params(nn), 1)

# output

1-element Vector{Float64}:
 0.9912108161055604

Keyword Arguments

• cse: perform common subexpression elimination when generating code (default true).

Symbolics stores an expression as a hash-consed directed acyclic graph but Symbolics.build_function prints it as a tree. Every time a subexpression is reused — the output of layer $n$ feeding each neuron of layer $n+1$, or the forward pass shared by every block of a symbolic gradient — the whole subtree is emitted again, so both the size of the generated code and the amount of redundant arithmetic grow exponentially with the depth of the network. With cse = true the graph is emitted as a let block of intermediate bindings instead, which keeps the code size proportional to the number of distinct nodes.

Pass cse = false to recover the fully inlined output; that is mostly useful for debugging, and for very small networks where the binding overhead is not amortised.

build_nn_function(eqs::AbstractArray{<:EquationSet}, sparams, svariables...)

Turn an array of equation sets into an executable function that returns an array of results.

Each entry of the array is built by the EquationSet method above, i.e. jointly; the entries themselves are independent of each other and stay separate functions.

Examples

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

c = Chain(Dense(2, 1, tanh))
nn = NeuralNetwork(c)
snn = SymbolicNeuralNetwork(nn)
eqs = [(a = c(snn.input, params(snn)),), (b = c(snn.input, params(snn)) .^ 3,)]
funcs = build_nn_function(eqs, params(snn), snn.input)
funcs([1.0, 2.0], params(nn))

# output

2-element Vector{NamedTuple{names, Tuple{Vector{Float64}}} where names}:
 (a = [0.985678060655224],)
 (b = [0.9576465981186686],)

build_nn_function(eq, nn)
build_nn_function(eq, nn, soutput)

Turn a symbolic equation into an executable function.

The result is called with the network input and the network parameters, and with the target output in between if the equation was built with one:

built_function(input, ps)
built_function(input, output, ps)

input may be a single sample (a vector), a batch (a matrix whose columns are the samples), or a batch with two batch dimensions (a three-dimensional array). See AbstractBatchedFunction for the shape of the result in each case.

build_nn_function(eq, sparams, svariables...)

The same, but with the symbolic parameters and the symbolic data variables given explicitly rather than taken from a SymbolicNeuralNetwork.

Keyword Arguments

• cse: perform common subexpression elimination when generating code (default true). See build_kernel.
• inplace: evaluate a batch with an in-place kernel (default true). See below.
• reduce: how to combine the results of the individual samples of a batch, either hcat (default) or +.

Examples

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

c = Chain(Dense(2, 1, tanh))
nn = NeuralNetwork(c)
snn = SymbolicNeuralNetwork(nn)
built_function = build_nn_function(c(snn.input, params(snn)), snn)
built_function([1.0, 2.0], params(nn)) ≈ c([1.0, 2.0], params(nn))

# output

true

Implementation

The equation is scalarised, a kernel that evaluates a single sample is generated from it (build_kernel or build_kernel!), and that kernel is wrapped in an AbstractBatchedFunction which adds the batching. Symbolics.build_function emits no in-place form for a scalar-valued equation, so those always take the out-of-place path.

build_nn_function(eqs::EquationSet, sparams, svariables...)

Turn a whole set of equations into one executable function, whose result has the same nesting as eqs.

Examples

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

c = Chain(Dense(2, 1, tanh))
nn = NeuralNetwork(c)
snn = SymbolicNeuralNetwork(nn)
eqs = (a = c(snn.input, params(snn)), b = c(snn.input, params(snn)) .^ 2)
funcs = build_nn_function(eqs, params(snn), snn.input)
funcs([1.0, 2.0], params(nn))

# output

(a = [0.985678060655224], b = [0.9715612392570434])

Implementation

All entries are generated as a single function whose flat result is split up again afterwards; see flatten_equations and split_result. Generating one function per entry instead would re-derive everything the entries have in common — for a symbolic gradient that is the whole forward pass, once per parameter array — and would compile one RuntimeGeneratedFunction per entry rather than one in total.

callee_name(e)

The name of the function a call expression calls, or nothing if e is not a call.

The emitted code refers to a function either by symbol, by function object, or by a qualified path (SymbolicUtils.Code.create_array), depending on how it was constructed. All three forms are reduced to a plain Symbol here so that the rules only have to deal with one of them.

carried_variables(layer)

Fresh symbolic arrays standing for whatever layer carries alongside the state at the seam, as a tuple. () by default, which is the plain-vector seam.

Extending

SymbolicNeuralNetworks.carried_variables(layer::MyLayer) = (Symbolics.variables(:c, 1:length(layer)),)

The arrays are built here rather than passed in because their shape is the layer's own knowledge — this package knows only input_dimension and output_dimension, which describe the state. Their name is the layer's own to choose too, and has to differ from the ones layer_seed uses itself — x for the state and λ for the sensitivities. Reusing one of those does not make a second array, it names the same one twice; layer_step refuses the seam rather than generating kernels that would quietly read the state where the carried datum should be.

Return () when there is nothing to carry, rather than an empty array: there would be nothing to hand the generated kernels, and an empty array is not a usable data argument in any case. The seam is then the plain vector it is for every other layer. SymplecticEuler over a system with no parameters is exactly this case.

The layer may still have to be given something — its functor takes the pair either way — in which case seam_value supplies it as a constant. Nothing varies, so nothing needs a variable, and the constant is folded into the expression like any other. See seam_interface.

checked_guess(loss, nn, ŷ, y; cse, inplace)

The guessed expression of loss as a function of prediction and target, or nothing if the guess cannot be trusted.

There are two ways for it not to be, and both mean the same thing to the caller — decline, and let SymbolicPullback fall back to the monolithic construction:

• the guess disagrees with loss, which is what represents_loss tests for;
• the guess cannot be built at all. A NetworkLoss need not accept a PassThroughLayer: the generic four-argument method of AbstractNeuralNetworks invites a loss to be written for the model it belongs to, and one written as (::MyLoss)(model::Chain, …) throws when passthrough_expression applies it to a model that is not a Chain. So does a model whose forward pass cannot be evaluated at the points represents_loss checks at.

The second case has to be caught here rather than left to the caller: layerwise = :auto promises a fallback, and a construction that throws instead of declining would break networks the monolithic path builds perfectly well.

Implementation

The try covers building and checking the guess, and nothing else. Once an expression is in hand and has been believed, differentiating it and generating code from it are this package's own work, and a failure there is a bug to surface rather than a reason to fall back.

checked_layer_seed(layer, key, prototype)

What layer_seed returns for layer, or nothing when the layer cannot be seeded at all.

A layer that carries something alongside the state and has not declared the seam interface (see seam_interface) has no seed, because the seam it is offered is a plain vector of symbolic variables. There are two ways for that to show, and both mean the same thing to the caller — decline, and let SymbolicPullback fall back to the monolithic construction:

• the layer returns more than the state, so state_expressions' default has no method for its output. This is GeometricMachineLearning's SymplecticEuler with return_parameters = true, which threads the parameters of the system on to the next layer and returns a Tuple;
• the layer cannot be applied to the bare vector at the seam in the first place, which is what the layer downstream of such a one does — its input is the tuple, and it has no method for anything else.

The second is why this catches rather than asking applicable(scalar_expressions, layer(sx, ps)): that predicate needs the layer to have been applied already, so it only covers the first.

Implementation

The try covers building the seed and nothing else — the same line checked_guess draws one level up. Once a seed is in hand, differentiating it and generating code from it are this package's own work, and a failure there is a bug to surface rather than a reason to fall back.

composes_layerwise(nn)

Whether composing the pullback layer by layer is the better choice for nn, which is what layerwise = :auto asks. True when the model decomposes into more than one step.

This is a question about which construction is preferable, not about whether the layerwise one applies: the layers still have to be seedable (checked_layer_seed) and the loss still has to reduce to a seed (loss_seed), both of which are settled afterwards by layerwise_gradient_function. So a true here can still be followed by a decline.

Implementation

A single layer is the one case where the monolithic construction is unambiguously right: there is no composition to keep out of the expression, so the two build the same derivative, and the seeded form of layer_step merely adds the sensitivity variables and a second generated function.

Above one layer the two are measured against each other by scripts/codegen_comparison.jl. In expression nodes — the size of the symbolic material each construction has to hold — and in seconds to build, for Dense chains with a FeedForwardLoss:

The layerwise node count is exactly linear — 864 more per identical added layer — against a column that multiplies by about six each time.

Build time crosses over at five layers, or at four of width eight. Below that the monolithic path is still ahead, by at most a fifth of a second, and this returns true there anyway: a threshold that reproduced the crossover would have to be a function of width as well as depth, fitted to timings from one machine, and it would be the wrong thing to get wrong. Choosing layerwise where monolithic would have been quicker costs a fixed fraction of a second; choosing monolithic on a network one layer deeper costs everything.

data_name(i)

The name the generated kernels give their i-th data argument: x1, x2, … There is no bound on how many there may be. One is the network input; a second is typically the target output of a loss; the layerwise pullback uses one per entry of a layer's seam plus one for the output sensitivities (see seam_interface).

decline(demanded, why)

Refuse the layerwise construction, why saying what stood in the way.

Returns nothing — the signal layerwise = :auto falls back to the monolithic construction on — or throws an ArgumentError naming why, when the caller asked for the construction by name with layerwise = true.

Implementation

Both outcomes go through here so that layerwise_gradient_function has one decision path: the message is raised where the decline happens rather than reconstructed afterwards by a second traversal of the same checks, which could not help but drift from them.

derivative(d)

The symbolic derivative stored in d.

derivative(g)

The symbolic gradient stored in g.

Examples

using SymbolicNeuralNetworks: SymbolicNeuralNetwork, Gradient, derivative, symbolic_parameter_gradient
using AbstractNeuralNetworks

c = Chain(Dense(2, 1, tanh))
nn = SymbolicNeuralNetwork(c)
g = Gradient(nn)

isequal(derivative(g), symbolic_parameter_gradient(g.f, nn))

# output

true

derivative(j)

The symbolic Jacobian stored in j.

evaluate_batch(f, data, ps)

Evaluate an AbstractBatchedFunction over the batch held in the tuple of data arguments data. There is one method per rank of the data arguments; the matrix one is the workhorse and the other two reduce to it.

flat_parameter_gradient(f, nn)
flat_parameter_gradient(f, sparams)

Differentiate the symbolic expression f with respect to the parameters of nn — or with respect to a set of symbolic parameters sparams given directly — as one flat object rather than as a set nested like the parameters.

A scalar f gives a vector of length flatlength(params(nn)); an array-valued f gives the $\mathrm{length}(f)\times\mathrm{flatlength}$ Jacobian

\[J_{ij} = \frac{\partial{}f_i}{\partial{}w_j},\]

with the rows indexed by vec(f) — the convention Jacobian uses for the derivative with respect to the input, and the matrix a Newton step is built from.

Pair it with build_flat_function for a function that is flat in both directions, and with NeuralNetworkParameters.unflatten to read a column block of the result back as the entry of the parameter set it belongs to.

The sparams form is the one for degrees of freedom that are not a network's parameters: nothing in either function reads a model, so an expression over a set of symbolic leaves goes through both. As with symbolic_parameter_gradient, which this calls, sparams may be a NetworkParameters or an EquationSet — the first is what a network's parameters are, the second what a caller writes out. build_flat_function, which this is paired with, is the narrower of the two and wants a NetworkParameters, because build_nn_function dispatches on one for its sparams.

Examples

The gradient of a scalar expression, against the same derivative in its nested form:

using SymbolicNeuralNetworks
using SymbolicNeuralNetworks: symbolic_parameter_gradient
using AbstractNeuralNetworks: Chain, Dense, params
using NeuralNetworkParameters: flatten

c = Chain(Dense(2, 1, tanh))
snn = SymbolicNeuralNetwork(c)
scalar = sum(c(snn.input, params(snn)))

flat = flat_parameter_gradient(scalar, snn)
nested, _ = flatten(symbolic_parameter_gradient(scalar, snn))
(length(flat), all(isequal.(flat, nested)))

# output

(3, true)

The Jacobian of a vector-valued one:

using SymbolicNeuralNetworks
using AbstractNeuralNetworks: Chain, Dense, params

c = Chain(Dense(2, 3, tanh))
snn = SymbolicNeuralNetwork(c)
size(flat_parameter_gradient(c(snn.input, params(snn)), snn))

# output

(3, 9)

flatten_equations(eqs)

Concatenate every entry of eqs into one vector of scalar equations, together with the NeuralNetworkParameters.ParameterLayout that records where each entry went.

Examples

using SymbolicNeuralNetworks: flatten_equations, SymbolicNeuralNetwork
using AbstractNeuralNetworks: Chain, Dense, params
using NeuralNetworkParameters: parameterrange

c = Chain(Dense(2, 3, tanh))
snn = SymbolicNeuralNetwork(c)
flat, layout = flatten_equations((a = c(snn.input, params(snn)), b = c(snn.input, params(snn)) .^ 2))
(length(flat), parameterrange(layout.children.a), parameterrange(layout.children.b))

# output

(6, 1:3, 4:6)

Implementation

The layout is the one NeuralNetworkParameters builds for a parameter set: an equation set has the same shape as one, its leaves are arrays (or single instances) of Num, and the layout records exactly what splitting the flat result needs — a range and a size per leaf, in the order the leaves are written. See unflatten_batch for the one thing that has to be added on top, and NeuralNetworkParameters.unflatten for the vector case, which needs nothing.

Each entry is normalised by scalar_expressions on the way in, which is what turns a Symbolics.Arr leaf into the Array{Num} the layout expects. The element type of the flat vector is fixed to Num rather than left to NeuralNetworkParameters.parameter_eltype to promote, so that it is the same type for every equation set — the code generation downstream dispatches on it.

flatten_gradient(gradient)

Lay a symbolic parameter derivative out flat: a single parameter-shaped set becomes a vector, and an array of them — which is what differentiating an array-valued expression gives, one set per entry — becomes a matrix with one row per entry. See flat_parameter_gradient.

function_arguments_and_body(expr)

Split a generated function expression into its vector of argument names and its body.

Throws an ArgumentError if expr is not of the shape Symbolics.build_function is documented to return, which turns an upstream change into an error at code-generation time rather than into subtly wrong code.

function_expression(argument_names, body)

Assemble an anonymous function definition; the inverse of function_arguments_and_body.

generated_expression(equation, svariables, sarrays; inplace, cse)

Call Symbolics.build_function and pick the half of its output that is asked for.

It returns an (out_of_place, in_place) pair for an array-valued equation and a single expression for a scalar-valued one; nothing is returned when the in-place half is asked for and there is none.

index_by_batch(expr, data_names)

Turn x[i] into x[i, k] for every data argument x.

The generated code reads a data argument as if it were a single sample. Adding the batch index k as a second index is what makes the same code read column k of a matrix instead, and is the reason the kernels take a batch index at all.

Both forms Symbolics emits for reading an entry are handled: x[i] (Expr(:ref, …)) and getindex(x, i), which it uses for arguments that were Symbolics.Arrs.

Throws an ArgumentError if a data argument is used other than by indexing it with a single index, because then the batch dimension would silently be ignored for that use.

Examples

using SymbolicNeuralNetworks: index_by_batch

index_by_batch(:(x1[1] + getindex(x1, 2)), (:x1,))

# output

:(x1[1, k] + getindex(x1, 2, k))

is_reserved_name(name)

Whether name is one a generated kernel gives an argument of its own, and therefore one a symbolic variable left free in an equation may not carry.

That is FIXED_NAMES together with the whole x1, x2, … family of data_names — the family and not just the arities in use, because whether a given name is generated would otherwise depend on how many data arguments the equation happens to have, and a free variable named x3 would pass the check today and break the day a third data argument arrived.

layer_seed(layer, key, prototype)

The scalar one layer's two derivatives are taken of, together with the variables it is written in: (seed, sparams, sdata, sλ), where

\[\mathrm{seed} = \lambda_k \cdot f_k(x_{k-1}; \theta_k).\]

sparams nests the layer's parameters under key, with the shape of prototype, so that the code generated from seed reaches into the parameter set of the whole network — see LayerStep. sdata is the tuple of data variables the seam is written in — the state first, then whatever carried_variables declared — and it and sλ are fresh for every layer, which is what keeps an expression built from this one dependent on that layer alone.

A layer that carries nothing gets sdata = (sx,), one plain vector, and this is then the construction it always was. A layer that carries something and has not declared the seam interface cannot be seeded at all — it either returns more than state_expressions' default can take apart, or has no method for a bare vector in the first place; checked_layer_seed is what turns that into a decline rather than an exception. See seam_interface.

Separate from layer_step so that scripts/codegen_comparison.jl measures the symbolic material this construction actually holds, rather than a second copy of it that can drift.

layer_step(layer, key, seeded; input_sensitivity, cse, inplace)

Build the LayerStep of one layer from seeded, the tuple layer_seed returns.

The seed is passed in rather than built here so that every layer of a chain can be seeded — and the chain declined if one of them cannot be, see checked_layer_seed — before any code is generated for any of them. Building a seed is one symbolic forward pass through one layer; generating its two kernels is the expensive half, and a chain that will be declined should not pay it.

A seam whose variables are not pairwise distinct is refused here rather than declined; see _assert_distinct_seam_variables.

input_sensitivity = false leaves out the derivative with respect to the layer's input, for the first layer of a chain, whose sensitivity is that of the loss to the network's input — something a parameter gradient has no use for, and which sweep accordingly never asks for. Generating it anyway would be half the code this function emits, spent on a function that is never called.

Implementation

Both derivatives come from differentiating the scalar

\[s_k = \lambda_k \cdot f_k(x_{k-1}; \theta_k),\]

in which $\lambda_k$ is a fresh vector of symbolic variables standing for the sensitivity of the loss to this layer's output. Its two gradients are precisely what the sweep needs: $\partial{}s_k/\partial{}x_{k-1}$ is $\lambda_{k-1}$ and $\partial{}s_k/\partial\theta_k$ is $\partial{}L/\partial\theta_k$.

Seeding the derivative like this rather than building $\partial{}x_k/\partial{}x_{k-1}$ and $\partial{}x_k/\partial\theta_k$ separately is what keeps both the expression and the generated code small: neither the Jacobian nor the rank-3 parameter derivative is ever materialised, and no contraction is left to do at run time. For four layers of width 16 the seeded form holds 68 760 nodes against 76 372 for the Jacobian-and-parameter-derivative pair.

The variables at the seam are fresh for every layer, which is the whole point — an expression built here refers to this layer's input and to nothing upstream of it, so its size depends on this layer alone. That the same names recur across layers is harmless, as each layer is compiled into its own kernel.

The two derivatives are compiled separately because they are reduced over a batch differently: a sensitivity is per-sample and concatenates, whereas the gradient of a batch is the sum of the per-sample gradients — which is what SymbolicPullback means by the pullback of a batch. So the sweep costs two calls per layer whatever the batch size, with no per-sample loop.

layerwise_gradient_function(nn, loss; demanded, cse, inplace)

Build the LayerwiseGradientFunction of loss for nn, or decline when the layerwise construction does not apply. There are three ways for it not to:

• the model does not decompose into steps (symbolic_steps);
• one of those steps cannot be seeded (checked_layer_seed);
• the loss cannot be reduced to a seed (loss_seed).

demanded = true — which is SymbolicPullback's layerwise = true — makes each of those an error naming the reason instead of the nothing that falls back.

Implementation

The layers are checked before the loss because the check is cheaper: loss_seed builds a function, evaluates it at three points and then differentiates and builds again, whereas seeding a layer is one symbolic forward pass. So a chain that cannot be seeded declines before any code is generated at all. The order also decides which reason a model that declines on both counts reports.

loss_expression(loss, ŷ, y)

The symbolic expression of loss as a function of the prediction ŷ and the target y, from which the layerwise pullback takes its seed $\partial{}L/\partial\hat{y}$.

Returns nothing by default, which means "not declared" — the package then guesses the expression with passthrough_expression and checks the guess before using it.

Extending

Declare the expression for a loss whose relation between prediction and target the guess cannot represent — one that carries extra data of its own, or that compares the prediction to the network's input rather than to output, as an autoencoder loss does:

SymbolicNeuralNetworks.loss_expression(loss::MyLoss, ŷ, y) = ...

A declared expression is used as given. It is deliberately not checked against loss the way the guess is: the reason to declare one is that the four-argument form means something the check would assume it does not, so checking it against that assumption would reject exactly the methods this exists for.

loss_seed(loss, model; cse, inplace)

Build the seed of the adjoint sweep: a function (ŷ, y, ps) -> ∂L/∂ŷ.

Returns nothing when the loss cannot be expressed as a function of prediction and target — which is what makes the layerwise construction fall back to the monolithic one instead of computing the wrong thing.

The expression comes from loss_expression if the loss declares one, and from checked_guess otherwise. Only the guess is checked against loss itself; a declared expression is used as given.

Implementation

ps is accepted and ignored: the expression has no parameters, and taking one anyway lets the seed be called exactly like the per-layer kernels of the sweep.

monolithic_gradient_function(nn, loss; cse, inplace)

Build the gradient of loss as one generated function, from one symbolic expression for the loss of the whole network.

This is what SymbolicPullback used to do unconditionally, and what it still does for a network the layerwise construction does not apply to, or is not worth applying to — see composes_layerwise.

Its cost is the reason for layerwise_gradient_function: the expression is O(width^depth) before anything is differentiated, and differentiating it walks the whole of it once per scalar parameter.

next_name!(counters, name)

Return the next unused name derived from name and count it in counters.

Examples

using SymbolicNeuralNetworks: next_name!

counters = Dict{Symbol, Int}()
(next_name!(counters, :var), next_name!(counters, :var))

# output

(:var_1, :var_2)

parameter_arguments(sparams)

Flatten a nested parameter set into the flat list of symbolic arrays that Symbolics.build_function is handed, together with the access path of each within the parameter object.

Symbolics.build_function only recognises a symbolic array that is passed to it as an argument, so passing the nested parameter object as a whole would leave its entries as free variables in the generated code. Flattening here and rebuilding the access paths in argument_substitutions keeps the kernel's own interface nested regardless.

Examples

using SymbolicNeuralNetworks: parameter_arguments, SymbolicNeuralNetwork
using AbstractNeuralNetworks: Chain, Dense, params

snn = SymbolicNeuralNetwork(Chain(Dense(2, 1, tanh)))
first(parameter_arguments(params(snn)))

# output

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

passthrough_expression(loss, ŷ, y)

Guess the expression of loss as a function of prediction and target, by applying it to a PassThroughLayer — a model whose prediction is its input.

This is right for every NetworkLoss that reaches its model exactly once, as model(input, ps), and compares the result to output, which is what the losses of AbstractNeuralNetworks do. It is wrong, rather than merely unavailable, for a loss that does something else: an autoencoder loss compares the prediction to the input, and so reads through a pass-through model as identically zero. That is why the guess is checked — see represents_loss — and why loss_expression exists.

postwalk(f, x)

Apply f to every node of the expression tree x, children first. Nodes that f returns are not visited again, so a rule may insert new syntax without it being rewritten in turn.

promoted_eltype(args...)

The element type the generated code will produce, promoted over all inputs.

This is needed because the in-place kernels write into an array that has to be allocated before they are called, so the element type cannot be taken from a result. Promoting over the inputs keeps Float32 parameters, symbolic (Num) inputs and ForwardDiff.Dual numbers working.

The walk over a nested parameter set is NeuralNetworkParameters.parameter_eltype, which promotes over the leaves and reaches the storage of a structured one — so a parameter that keeps fewer numbers than its interface shows contributes the element type of the numbers it actually keeps.

Note that this derives the element type from the inputs, not from the expression, which the out-of-place path does instead. The two can differ: an equation over integer inputs and integer parameters evaluates to a Float64, which no Array{Int} can hold, so the allocators widen an integer type with float. A Float32 network whose generated code contains a Float64 constant is rounded to Float32 rather than widened — which is the behaviour one wants for the network, but is worth being aware of.

reference_parameters(nn)

A numeric parameter set with the shape of nn's symbolic one, filled deterministically.

Implementation

The shape comes from the symbolic parameters through their NeuralNetworkParameters.ParameterLayout, and the numbers from unflattening a flat vector of them: a layout built over Num leaves unflattens a Float64 vector into Float64 leaves, since a leaf is rebuilt from its prototype's shape and not from its element type. Going through the layout rather than through initialparameters is what keeps this free of the global RNG, and it gets structured parameters right for free.

represents_loss(loss, nn, value)

Whether the built passthrough_expression value agrees with loss itself, on nn's model, at a handful of points.

true when loss cannot be evaluated in the four-argument numeric form at all, in which case there is nothing to compare against and the expression is taken at its word — that is the case an overriding method for a loss with its own call signature lands in.

Implementation

Neither the parameters (reference_parameters) nor the points are random: whether a pullback builds must not depend on the state of the global RNG, and building one must not advance it — a caller who seeds the RNG and then builds a pullback would otherwise get different data afterwards than before.

Three points, so that an expression which happens to agree at one of them is still rejected. The case in mind is an autoencoder loss, which compares the prediction to the input and so reads as identically zero through a pass-through model — that agrees with the real loss exactly when the real loss is zero too.

The three differ in direction and not merely in scale: three points on one ray through the origin would leave an expression that is right along that ray and wrong everywhere else undetected, which is no more work to avoid than to allow.

scalar_expressions(eq)

Normalise an equation to scalar expressions, i.e. to a Num or an Array{Num}.

A Symbolics.Arr — the type @variables x[1:n] produces — is a symbolic object in its own right whose entries are only materialised by Symbolics.scalarize. Symbolics.build_function cannot generate code for one, so every equation passes through here on its way into build_nn_function. Arrays of Num (what this package builds, see symbolic_variables) pass through unchanged.

seam_arguments(layer, x)

The data arguments the kernels generated for layer are called with, given the layer's run-time input x: a tuple in the order (state, carried...) that carried_variables declared. (x,) by default.

Extending

SymbolicNeuralNetworks.seam_arguments(::MyLayer, x::Tuple) = (first(x), flat(last(x)))

This is the run-time counterpart of seam_value, and the two have to agree: what the seam is written in and what it is called with are the same list. Every entry must have the same rank and batch size as the state. See seam_interface.

The seam interface

layer_seed puts fresh symbolic variables between two layers, which is what keeps the symbolic material a sum over layers rather than a product. By default those variables are one plain vector — the layer's state — because that is what a Dense maps to a Dense.

A layer may carry more than the state. GeometricMachineLearning's SymplecticEuler threads the parameters of the system through the chain alongside it, so it takes and returns a pair. Four functions say how such a layer meets the seam; each defaults to exactly the plain-vector construction, so a layer that carries nothing needs none of them, and a layer that carries something declares all four together:

The carried data is data, never a differentiation target: $\lambda$ pairs with the state, the seed is differentiated with respect to the state and with respect to the layer's parameters, and the carried variables are extra arguments of the generated kernels. So a layer that carries something gets the same two kernels as any other, taking one more argument each.

seam_arguments must return arrays with the same rank and batch size as the state, in the order (state, carried…) that carried_variables declared — the constraint every generated function of this package imposes on its data arguments. For carried data that is the same for the whole batch that means broadcasting it out to one column per sample.

carried_variables must return variables of its own: layer_seed names the state x and the sensitivities $\lambda$, and a carried array that reuses either name is the same symbolic array rather than a second one. layer_step rejects that rather than letting it through, because nothing downstream would — Symbolics.build_function binds one array to two argument slots and the generated code reads both from the last one, which builds and runs and returns a wrong gradient.

seam_value(layer, sx, carried...)

What layer is applied to at the seam, assembled from the state variables sx and the arrays carried_variables returned. sx alone by default.

Extending

SymbolicNeuralNetworks.seam_value(layer::MyLayer, sx, sc) = (sx, sc)

A layer whose carried data reaches it in some other shape than the flat array the kernels take — SymplecticEuler wants a NamedTuple of system parameters — reassembles it here, and takes it apart again in seam_arguments. See seam_interface.

split_result(layout, out)

Split the flat result out of a jointly generated function into the nesting recorded in layout.

out is dispatched on by its number of dimensions, which is how the layout of the batch is accounted for — see AbstractBatchedFunction for where those layouts come from. A vector, which is what a summed batch or a single sample produces, is the case NeuralNetworkParameters.unflatten already covers: every entry simply keeps the shape of its equation. Anything else has a batch dimension and goes to unflatten_batch.

state_expressions(layer, y)

The part of layer's output $\lambda$ pairs with, as scalar expressions. All of it by default.

Extending

SymbolicNeuralNetworks.state_expressions(::MyLayer, y) =
    SymbolicNeuralNetworks.scalar_expressions(first(y))

A layer that passes its carried data on returns it beside the state, and the seed is $\lambda_k\cdot{}f_k(x_{k-1};\theta_k)$ over the state — so this is what says which part that is. See seam_interface.

step_parameters(step, ps)

The entry of the parameter set ps that belongs to step.

structured_parameters(f, w)

The parameter set w stands for: w laid out according to f's layout, or — when w carries a layout of its own — according to that one.

substitute_symbols(expr, substitutions)

Replace every symbol of expr that is a key of substitutions by the corresponding value.

sweep(steps, x, ps, seed, output)

The derivative of the loss with respect to each step's parameters, as a tuple in the order of steps.

Implementation

The recursion runs the forward pass on the way in and the adjoint sweep on the way out, so each intermediate result lives on the stack for exactly as long as the backward pass needs it, and the whole sweep stays type stable over the tuple of steps.

The first step's dλ is never called: it would give the sensitivity of the loss to the network's input, which a parameter gradient has no use for. Hence the two entry points — adjoint_step does the general case, and this function drops that one call. The first step is not built with a dλ at all, so dropping it here is what makes that legal as well as cheaper; see layer_step.

symbolic_derivative(f, differentials)

Differentiate the scalar expression f with the differential operators in differentials, keeping their shape and nesting. Together with symbolic_differentials this is what turns "the parameters of a network" into "the derivative of f with respect to each of them".

symbolic_differentials(svariables)

The differential operators belonging to a set of symbolic variables, with the same shape and nesting as the variables themselves. See symbolic_derivative.

symbolic_parameter_gradient(f, nn)
symbolic_parameter_gradient(f, sparams)

Differentiate the symbolic expression f with respect to the parameters of nn, or with respect to a set of symbolic parameters sparams given directly.

The result has the same nesting as the parameters of nn. For an array-valued f it is an array of such parameter sets, one per entry of f.

The second form is what the degrees of freedom of an expression that is not a network's forward pass go through: nothing here reads the model, only the parameters. sparams is anything symbolic_differentials can walk: a NetworkParameters, which is what the parameters of a SymbolicNeuralNetwork are, or an EquationSet of symbolic leaves.

This is used by Gradient and by SymbolicPullback.

Examples

using SymbolicNeuralNetworks: SymbolicNeuralNetwork, symbolic_parameter_gradient
using AbstractNeuralNetworks
using AbstractNeuralNetworks: params

c = Chain(Dense(2, 1, tanh))
nn = SymbolicNeuralNetwork(c)
symbolic_parameter_gradient(c(nn.input, params(nn)), nn)[1].L1.b

# output

1-element Vector{Symbolics.Num}:
 1 - (tanh(W_2₁ + W_1₁ˏ₁*x₁ + W_1₁ˏ₂*x₂)^2)

symbolic_steps(nn)

The sequence of steps nn's model decomposes into, as (layer, key) pairs, or nothing when it does not decompose into one.

A model decomposes when it is a Chain whose layers correspond one-to-one to the entries of the parameter set, and each of whose layers knows the dimensions it maps between — which is what the layerwise pullback needs in order to put fresh variables at the seams. Anything else takes the monolithic path.

symbolic_variables!(counters, x, name)

The workhorse of symbolic_variables. counters maps a name to the number of variables that have been created under it so far and is updated in place, which is what makes the leaves of a nested parameter set distinguishable.

symbolic_variables(x, name)

Build symbolic variables with the shape of x, named after name.

x may be a number, an array, or an arbitrarily nested NamedTuple/NetworkParameters of those — i.e. anything that can hold the parameters of a neural network. Every leaf gets its own name, numbered consecutively: name_1, name_2, …

Examples

using SymbolicNeuralNetworks: symbolic_variables

symbolic_variables((a = 1.0, b = [1, 2]), :X)

# output

(a = X_1, b = Symbolics.Num[X_2₁, X_2₂])

using SymbolicNeuralNetworks: symbolic_variables
using AbstractNeuralNetworks: NeuralNetwork, Chain, Dense, params
using NeuralNetworkParameters: NetworkParameters

nn = NeuralNetwork(Chain(Dense(1, 2; use_bias = false), Dense(2, 1; use_bias = false)))
sparams = symbolic_variables(params(nn), :W)
(sparams isa NetworkParameters, keys(sparams), size(sparams.L1.W), eltype(sparams.L1.W))

# output

(true, (:L1, :L2), (2, 1), Symbolics.Num)

Implementation

The variables are scalar ones (Symbolics.variable/Symbolics.variables), so an array of parameters becomes an Array{Num} rather than a Symbolics.Arr. Symbolics cannot differentiate with respect to the entries of a Symbolics.Arr without scalarising it first, and Symbolics.build_function cannot generate code for expressions that still contain one; using scalar variables throughout avoids both problems. See scalar_expressions.

trailing_dimensions(equation_size)

The number of columns one sample of an equation of size equation_size occupies, i.e. the product of all but its first dimension. A scalar-valued equation counts as one of size $m = 1$, so its empty size gives 1 — Base.tail would throw on it.

two_batch_dimension_message(equation_size)

The error text for an equation whose result already uses the second dimension. Shared with unflatten_batch, so that an entry of an equation set is rejected in the same words as the same equation built on its own.

unflatten_batch(layout, out)

Split a batched flat result into the nesting recorded in layout, giving each entry the shape AbstractBatchedFunction documents for a concatenated batch:

• a $P\times{}N$ matrix, in which an entry of size $(m, n, \ldots)$ becomes an $m\times(n\cdot\ldots\cdot{}N)$ matrix,
• a $P\times{}N_1\times{}N_2$ array, in which it becomes an $m\times{}N_1\times{}N_2$ one.

A scalar-valued entry is treated as one of size $m = 1$ throughout, so it comes back as a $1\times{}N$ matrix.

This is deliberately not a method of NeuralNetworkParameters.unflatten, which already means something else for a matrix: splitting the rows of a Jacobian taken with respect to a flat parameter vector, with no batch dimension to restore.

Implementation

Each entry is copied out of out rather than viewed into it, so that the entries are ordinary Arrays and cannot alias each other.

unseedable_reason(steps, seeds)

The why decline is given when one of a chain's layers cannot be seeded: the keys and types of every layer whose entry in seeds came back nothing.

use_base_mapreduce(expr)

Replace Symbolics._mapreduce with Base.mapreduce.

Symbolics._mapreduce cannot be differentiated by Zygote, whereas Base.mapreduce can. Its trailing Colon(), (:init => false,) arguments are the positional form of dims = Colon().

Nothing this package generates contains a Symbolics._mapreduce any more — that used to come from reductions over un-scalarised Symbolics.Arrs, which scalar_expressions now rules out. The rule is kept because the equations a user passes in are not under our control.

Examples

using SymbolicNeuralNetworks: use_base_mapreduce

use_base_mapreduce(:(Symbolics._mapreduce(identity, +, x, Colon(), (:init => false,))))

# output

:(mapreduce(identity, +, x; dims = Colon()))

use_generic_array_constructor(expr)

Replace SymbolicUtils.Code.create_array(typeof(…), …) with create_array(Array, …).

create_array takes the array type to construct as its first argument, and Symbolics fills that in with the type of one of the arguments of the generated function. For us that argument is the parameter set — a NamedTuple, from which no array can be constructed — and even where it is an array it may be a SubArray or a ReshapedArray that create_array has no method for. Array is the generic choice that works in every case.

Examples

using SymbolicNeuralNetworks: use_generic_array_constructor

use_generic_array_constructor(:((SymbolicUtils.Code.create_array)(typeof(ps), nothing, Val{1}(), a)))

# output

:(SymbolicUtils.Code.create_array(Array, nothing, Val{1}(), a))