Walking a parameter set

mapparameters(f, ps)
mapparameters(f, ps, rest...)

Apply f to every leaf of ps, returning a parameter set of the same shape.

With further arguments the trees are walked in lockstep and f receives one leaf from each, which is how two parameter sets are combined entrywise. Their keys have to agree.

f sees whole leaves — a SymmetricMatrix arrives as a SymmetricMatrix. Use mapstorage to see only the differentiable storage instead.

Examples

using NeuralNetworkParameters

ps = NetworkParameters((L1 = (W = [1.0 2.0], b = [3.0]),))
mapparameters(x -> 2x, ps).L1.W

# output

1×2 Matrix{Float64}:
 2.0  4.0

Combining two sets:

using NeuralNetworkParameters

a = NetworkParameters((L1 = (b = [1.0, 2.0],),))
c = NetworkParameters((L1 = (b = [10.0, 20.0],),))
mapparameters(+, a, c).L1.b

# output

2-element Vector{Float64}:
 11.0
 22.0

mapstorage(f, ps, rest...)

Like mapparameters, but f is applied to the freeparameters of each leaf and the leaf is rebuilded around the result.

This is the level at which entrywise arithmetic on a structured parameter is meaningful. Halving a SymmetricMatrix means halving the $n(n+1)/2$ numbers it stores; broadcasting over its dense $n \times n$ interface would do twice the work, and for a skew-symmetric or triangular matrix there is no setindex! to broadcast through at all.

A nothing in place of a leaf reaches f as far as there is one leaf to pair it with: the storage of a SymmetricMatrix or a manifold element is a single array, so f is handed that array and the nothing. A leaf whose storage is several blocks has nothing to pair one nothing with, and raises as a nothing branch does — an out-of-place walk has nothing to put in the hole either way.

Examples

using NeuralNetworkParameters

ps = NetworkParameters((L1 = (W = [1.0 2.0; 3.0 4.0],),))
mapstorage(x -> x ./ 2, ps).L1.W

# output

2×2 Matrix{Float64}:
 0.5  1.0
 1.5  2.0

mapparameters!(f, dest, srcs...)

Walk dest and srcs in lockstep, calling f(dest_leaf, src_leaves...) for its effect on dest_leaf, and return dest.

This is the walk an optimizer runs every iteration, and it allocates nothing at any width or depth. As with foreachparameters, a nothing in srcs skips that position, and the keys of a keyed branch have to agree.

mapstorage!(f, dest, srcs...)

Like mapparameters!, but f is handed the freeparameters of each leaf.

No rebuild is needed: the storage of a leaf is the leaf's own memory, so writing into it is writing into the leaf. Allocation-free on the same terms.

foreachparameters(f, ps, rest...)

Walk the leaves of ps for the side effect of f, in lockstep with rest. Returns nothing.

The children of a keyed branch are paired by key, and the keys have to agree — a rest whose keys are the same set in a different order is an ArgumentError and not a silent crossing-over. A Tuple branch is paired positionally, since the blocks of a multi-block leaf have no keys to agree on.

A branch or leaf of rest that is nothing skips that position entirely — f is not called there. This is what lets a gradient tree that is missing the entries of a frozen or non-trainable layer be walked against the parameters it belongs to, without having to fill the holes in first.

Allocation-free, at any width of branch, any depth of nesting and any number of rest: the branches are indexed in place rather than taken apart, so nothing is materialised on the way in.

See mapparameters! for the in-place variant that returns its destination.

foldparameters(op, init, ps)
foldparameters(op, init, ps, rest...)

Left-fold op over the leaves of ps, starting from init.

With further arguments the trees are walked in lockstep and op receives one leaf from each — op(acc, a_leaf, b_leaf, …) — which is how an inner product or a quadrature norm over a parameter set is taken without flattening it first. The children of a keyed branch are paired by key and the keys have to agree, exactly as for mapparameters; a Tuple branch is paired positionally, since the blocks of a multi-block leaf have no keys to agree on.

The leaves are visited in the order flatten writes them, so a fold and a flattening agree.

op sees whole leaves — a SymmetricMatrix arrives as a SymmetricMatrix. Use foldstorage to fold over the differentiable storage instead.

Where foreachparameters skips a set that is nothing, a fold raises: it reduces every leaf it is given, so a set left out would make the result a partial sum without saying so. A nothing in place of a single leaf still reaches op, exactly as it reaches f in mapparameters — what a missing leaf contributes to the sum is the caller's to decide, and only the caller's.

Allocation-free at any width, depth and arity — provided a caller that hands op on through a function of its own annotates it ::F where {F} there. Julia does not specialise on a function argument it never sees called, so without that annotation op arrives boxed and every leaf costs a dynamic dispatch: 3 088 bytes a call on a 369-leaf set at arity one and 6 160 at arity two, against zero with it, identically on Julia 1.11.9 and 1.13.0-rc3. A closure that captures a function needs nothing, a closure being its own type, which is why (acc, x) -> acc + abs2(f(x)) is the way to fold a function of each leaf and no second function argument is needed here.

Examples

using NeuralNetworkParameters

ps = NetworkParameters((L1 = (W = [1.0 2.0], b = [3.0]), L2 = (W = [4.0;;],)))
foldparameters((n, x) -> n + length(x), 0, ps)

# output

4

Two sets in lockstep, which is $\sum_i a_ib_i$ without a flat vector of either:

using NeuralNetworkParameters

a = NetworkParameters((L1 = (W = [1.0 2.0], b = [3.0]),))
b = NetworkParameters((L1 = (W = [4.0 5.0], b = [6.0]),))
foldparameters((acc, x, y) -> acc + sum(x .* y), 0.0, a, b)

# output

32.0

foldstorage(op, init, ps, rest...)

Like foldparameters, but op is handed the freeparameters of each leaf.

This is the level at which a reduction over a structured parameter is meaningful, and it is the level flatten writes. The pairing of a SymmetricMatrix is over the $n(n+1)/2$ numbers it stores; reading its dense $n \times n$ interface instead would count every off-diagonal entry twice, and for a skew-symmetric or triangular matrix there is no dense reading to be had at all.

A nothing in place of a leaf reaches op as far as there is one leaf to pair it with, on the terms mapstorage states: the storage of a SymmetricMatrix is a single array, and a leaf whose storage is several blocks raises instead, since one nothing cannot stand for each of them.

Examples

using NeuralNetworkParameters

a = NetworkParameters((L1 = (b = [1.0, 2.0],),))
b = NetworkParameters((L1 = (b = [3.0, 4.0],),))
foldstorage((acc, x, y) -> acc + sum(x .* y), 0.0, a, b)

# output

11.0