Converting between them

ParameterLayout

Where each leaf of a parameter set lands in its flat vector, and what has to be done to put it back.

A layout is an ordinary value, built once by parameterlayout and then reusable: it can be stored in an optimizer cache, handed to a solver, or compared with ==. The alternative — returning a closure that undoes the flattening, as ParameterHandling.flatten does — is none of those things, and a chain of closures is not type stable either.

The five concrete layouts mirror the five things a parameter set is made of: ParametersLayout for a NetworkParameters, NestedLayout for a NamedTuple, TupleLayout for a Tuple, WrappedLayout for a leaf with structured storage, and LeafLayout for one whose numbers are copied directly.

parameterlayout(ps)

Build the ParameterLayout of ps in one walk.

The leaves are laid out in the order they are encountered, depth first, so the flat vector reads like the parameter set does.

Examples

using NeuralNetworkParameters

ps = NetworkParameters((L1 = (W = [1.0 2.0], b = [3.0]), L2 = (W = [4.0;;],)))
layout = parameterlayout(ps)
(length(layout), parameterrange(layout.inner.children.L1.children.b))

# output

(4, 3:3)

parameterrange(layout)

The stretch of the flat vector that layout occupies.

flatlength(ps)

The number of entries ps flattens to, without allocating the flat vector.

Deliberately not called parameterlength: AbstractNeuralNetworks has a function of that name for the parameter count of a model, and the two would collide on using.

This returns an Int, so nothing downstream of the call depends on the layout's type. Anything that wants only the size should call it; anything that has to unflatten later needs the layout and cannot.

That used to be a large difference in compile time as well — on Julia 1.11, 1.26 s here against 13.20 s through parameterlayout on a 369-leaf NetworkParameters. It is not one any more: 0.2.3 took that set to 1.05 s through parameterlayout against 1.03 s here. The reason to prefer this is the type it hands back, not the clock.

flatten(ps)
flatten(T, ps)

Copy every number of ps into one flat Vector, and return it together with the ParameterLayout needed to put it back.

Without T the element type is parameter_eltype(ps), i.e. the parameters' own — a Float32 network flattens to a Vector{Float32}.

Examples

using NeuralNetworkParameters

ps = NetworkParameters((L1 = (W = [1.0 2.0], b = [3.0]),))
v, layout = flatten(ps)
v

# output

3-element Vector{Float64}:
 1.0
 2.0
 3.0

unflatten is its inverse:

using NeuralNetworkParameters

ps = NetworkParameters((L1 = (W = [1.0 2.0], b = [3.0]),))
v, layout = flatten(ps)
unflatten(layout, v) == ps

# output

true

Implementation

The copy is a copyto! per leaf over a known range, so it runs at memory bandwidth and works unchanged for GPU arrays — no element is ever indexed individually.

Copying rather than viewing is deliberate. A flat vector of views could not carry a different element type from the parameters, and that is exactly what differentiating through the flat form needs: the unflatten on the forward pass of ForwardDiff has to produce Dual-valued parameters over a Dual-valued vector. See flatten! for the allocation-free form used in inner loops.

flatten!(v, ps, [layout])

Write the numbers of ps into the existing vector v, and return v.

Allocation-free when the layout is supplied, which is the point: an optimizer that flattens its parameters once per step, or twice per inner product, should not allocate a fresh vector each time.

v, layout = flatten(ps)          # once
flatten!(v, ps, layout)          # per iteration, zero allocations

unflatten(layout, J::AbstractMatrix)

Split the rows of J into the shape of the parameter set — for a Jacobian taken with respect to the flat vector, so that the block belonging to each parameter can be read off.

Each leaf becomes the $n_\mathrm{leaf} \times \mathrm{size}(J, 2)$ row block it occupies. The leaves are not rebuilt: a block of a Jacobian is not a parameter, so there is nothing to rebuild it into.

unflatten(fp::FlatParameters)

The structured form of fp, using the layout it carries.

unflatten!(ps, layout, v)

Write the numbers of v into the leaves of the existing ps, and return ps.

Allocation-free on the same terms as flatten!, whose counterpart it is. The write goes through freeparameters, so only the storage of a structured leaf is touched — a SymmetricMatrix has its n(n+1)/2 numbers replaced and stays symmetric.

Requires mutable leaves; a parameter set with a scalar leaf has to use the out-of-place unflatten.