The leaf protocol

freeparameters(x)

The differentiable storage of a leaf x — the numbers that a flat parameter vector should contain, and the coordinates an optimizer should work in.

For an ordinary array this is the array itself. For a structured type it is whatever the type keeps its degrees of freedom in: an $n \times n$ symmetric matrix stores $n(n+1)/2$ numbers, and it is those that belong in the flat vector — not the $n^2$ entries of the dense interface, which do not even have the right length.

The return value may be

• the leaf itself, which marks it as terminal — the recursion stops and the numbers are copied straight out of it;
• another array or number;
• a Tuple or NamedTuple of either, for a type whose freedom lives in several blocks.

Together with rebuild this is the whole extension protocol: define the two for a type and it can be flattened, walked and written to HDF5, with nothing in this package knowing about it.

Extending

NeuralNetworkParameters.freeparameters(A::SymmetricMatrix) = A.S
NeuralNetworkParameters.rebuild(A::SymmetricMatrix, data)  = SymmetricMatrix(data, A.n)

Those two are the whole protocol for a leaf that is an AbstractArray subtype, which every structured parameter type in the ecosystem is. A leaf that keeps its numbers behind some other interface needs one method more, because parameter_eltype must not raise and so cannot ask an arbitrary type where its storage is:

NeuralNetworkParameters.parameter_eltype(x::MyLeaf) = parameter_eltype(freeparameters(x))

Without it the leaf contributes nothing to the promoted element type, and flatten says so rather than guessing one.

GeometricOptimizers already exposes exactly this relation as Base.parent for its manifolds, VectorStorageMatrixes and horizontal lifts, so one delegating method covers all of them:

NeuralNetworkParameters.freeparameters(x::Union{Manifold, VectorStorageMatrix, AbstractLieAlgHorMatrix}) =
    parent(x)

Base.parent is deliberately not used as the protocol here: parent(::SubArray) is the whole underlying buffer rather than the view's own entries, so Base's relation and this one are not the same.

Examples

using NeuralNetworkParameters: freeparameters

A = [1.0 2.0; 3.0 4.0]
freeparameters(A) === A

# output

true

rebuild(prototype, data)

Rebuild a leaf shaped like prototype from the storage data. The inverse of freeparameters, and for an ordinary array simply data.

prototype is the leaf the storage came from, which is what carries the non-differentiable information along: the n of a SymmetricMatrix, the N and n of a horizontal lift. Taking them from a prototype rather than from the type is what lets data have a different element type from prototype — ForwardDiff.Duals, when a flattened parameter set is differentiated through — and it is also why the concrete type comes back unchanged, where reconstructing from a type name can quietly return a sibling type instead.

See freeparameters for how to extend this.

parameter_eltype(ps)

The element type to flatten ps into: promote_type over the element types of its leaves.

Float32 parameters therefore flatten to a Vector{Float32}. Defaulting to Float64 instead — as ParameterHandling.flatten does — silently doubles the width of every single-precision network that passes through, and every quantity computed from the flat vector downstream with it.

Examples

using NeuralNetworkParameters: parameter_eltype

parameter_eltype((a = Float32[1, 2], b = Float32[3;;]))

# output

Float32

A mixed parameter set promotes:

using NeuralNetworkParameters: parameter_eltype

parameter_eltype((a = Float32[1, 2], b = [3.0]))

# output

Float64

Implementation

This function is total, where freeparameters is not: every NetworkParameters runs its constructor through here, including sets that hold no numbers at all — a gradient tree with nothing where an untouched layer's entries would be, or SymbolicNeuralNetworks wrapping generated functions in one. A leaf this package cannot read numbers out of contributes nothing to the promotion, exactly as an empty set does, and both report Union{}. A set that cannot be flattened is still a set with an element type; the protocol error comes from parameterlayout, which decides something with it.

The recursion follows freeparameters for an AbstractArray leaf — every structured parameter type in the ecosystem — and a leaf that keeps its numbers behind another interface opts in with one method more, described under freeparameters. Asking an arbitrary type where its storage is would mean raising, which this function must not do.

It follows it at the level of values, freeparameters(x) === x, and that is a guarantee rather than an implementation detail. A structured leaf reports the element type of its storage, which need not be the eltype of the interface it presents: a leaf that is an AbstractMatrix{Float64} over a Vector{Float32} flattens into a Vector{Float32}, because that is what its numbers are. Deciding the promotion from the leaf's type instead would read the interface and be wrong about such a leaf, which is why it is not decided there — and there is no cost to buy with it.

For a NetworkParameters the answer is already on the type, put there by its constructor, so nothing is recomputed and the call folds to a constant. A bare NamedTuple or Tuple is read in place at literal indices, one @generated body per branch shape, for the reason the head of walk.jl gives — so it too costs nothing at any width, depth or shape, and flatten(ps) costs exactly what flatten(T, ps) costs.