Take the remaining parameter walks from the packages that own them

CHANGELOG.md

@@ -9,10 +9,136 @@ breaking release).
 > [!NOTE]
 > Entries for 0.1.0 through 0.4.8 were reconstructed from git history, the release tags and the
 > merged pull requests, not written at the time. They are accurate about *what* changed and are
-> deliberately coarser about detail than the 0.5.0 and 0.6.0 sections below, which were written
+> deliberately coarser about detail than the 0.5.0 and later sections below, which were written
 > alongside the work. Where a release removed exported names the list is given; where it is a
 > reconstruction of intent, it says so.
 
+## [0.7.0]
+
+**The traversal of a parameter set now belongs to the package that owns the parameters, and the
+traversal of a `NamedTuple` belongs to `Base`.** 0.6.0 handed the HDF5 walk over to
+[NeuralNetworkParameters.jl][nnp] and `GeometricOptimizers`; this release finishes the job for the
+remaining walks. `map_to_cpu` becomes one walk, `apply_toNT` turns out to have been `Base.map` all
+along, and `_eltype` turns out to have been a hand-rolled `parameter_eltype`.
+
+**It also makes the optimizer cache immune to a change coming in `GeometricOptimizers`**, which is
+the half of this release with no visible effect today — see *Fixed*.
+
+### Removed (breaking)
+
+- **`apply_toNT` is gone from the export list and from the package. It was `Base.map`.**
+
+  ```julia
+  apply_toNT(f, a, b)   →   map(f, a, b)
+  ```
+
+  `map` over `NamedTuple`s takes any number of arguments and already throws
+  `ArgumentError: Named tuple names do not match.` on mismatched *or* reordered keys — which is what
+  the hand-rolled `@assert keys(ps[1]) == keys(p)` was approximating, except that Base's check cannot
+  be compiled out the way an `@assert` can. Heterogeneous values map fine, so a `StiefelManifold`
+  beside an ordinary `Matrix` is no obstacle. Verified on Julia 1.10, the compat floor, as well as on
+  1.13.
+
+  `_norm`, `_diff` and `_add` use `map` directly; that is the faithful translation, not a
+  simplification, because `_diff` and `_add` recurse through their own `NamedTuple` methods and
+  `_norm` divides by `√length` one level down. `GeometricOptimizers` carried a
+  character-identical copy of the same function, reached from here by qualified call; that copy goes
+  in its own release, and this change is what frees it.
+
+- **`_eltype` is gone; `NeuralNetworkParameters.parameter_eltype` replaces it.** The two are not the
+  same function — `_eltype` returned the element type of the *first* leaf and read a structured leaf's
+  dense interface, where `parameter_eltype` promotes across every leaf and descends through
+  `freeparameters` — but at the four call sites this package had they cannot disagree, and it is worth
+  saying why rather than claiming a fix that could not fire.
+
+  `_eltype` was only ever asked for a `T` in two places: under `_use_go_cache`, which requires
+  `x isa GeometricOptimizers.OptimizerSolution`, and on the `ps_leaf` that reaches
+  `_leaf_optim_step!`, which is such an `x`. And `OptimizerSolution{T}` is homogeneous in `T` by
+  construction — its `NamedTuple` arm is
+  `ArrayNamedTuple{T} = NamedTuple{S,<:Tuple{Vararg{AbstractArray{T}}}}`. A layer mixing `Float32`
+  and `Float64` weights therefore *fails* that test and recurses to one cache per weight, so a
+  first-leaf answer and a promoted one were never different answers. What the substitution buys is
+  four fewer methods to own and the upstream spelling at the point where the walk is upstream's.
+  Unexported, so this is breaking only for code reaching into the package.
+
+- **`add!(::NamedTuple, ::NamedTuple, ::NamedTuple)` is gone, and `_add` and `add!` are gone from the
+  export list.** The container arm of `add!` had no caller in the package, the tests, the docs or the
+  scripts, and `AbstractNeuralNetworks.add!` — whose generic it was a method of — is about a
+  destination and two summands, which a parameter *tree* is not. `add!` remains available from the
+  package that owns the generic, this one only adding methods for the structured matrix types:
+
+  ```julia
+  using AbstractNeuralNetworks: add!
+  ```
+
+  `_add`'s two siblings `_norm` and `_diff` were never exported, and they are the two of the three
+  that anything in `src/` actually calls; `_add` was the odd one out. Qualified, it still works.
+
+- **`_add(::History, ::SingleHistory)` is now `_push_history!`.** Unexported and internal, one caller.
+  Two unrelated meanings on one name is one too many, and the new name says that it mutates its first
+  argument, which the old one hid.
+
+### Changed
+
+- **`map_to_cpu` is one walk instead of eight methods.** `NeuralNetworkParameters.mapstorage` hands a
+  function the storage of a leaf and rebuilds the leaf around the result, so the five methods that
+  existed to unwrap and reconstruct a `StiefelManifold`, a `SymmetricMatrix`, a `SkewSymMatrix` and
+  the two triangular types collapse into one, the two that recursed into a `NetworkParameters` and a
+  layer go with them, and the plain-array one is all that is left — as `_to_host`, the function handed
+  to the walk. `GeometricOptimizers` supplies the protocol for its own types, so nothing here knows
+  which structured types exist and one added upstream is covered without a change on this side.
+
+  `mapstorage` and not `mapparameters`: the latter hands the function *whole* leaves, which would
+  still need a method per type to reach the storage. The `NeuralNetwork` method stays — the docs
+  tutorials and several scripts call it on a whole network.
+
+  It was **untested**, which is why the rewrite comes with `test/map_to_cpu_tests.jl`: that every
+  structured leaf comes back as the type it went in as, that the `n` a structured leaf carries
+  survives although it is not in its storage, that element types are preserved, that the leaves are
+  copies rather than the same arrays, and that a whole network keeps its architecture, model and
+  backend.
+
+### Fixed
+
+- **The shape of the optimizer cache no longer depends on which types `GeometricOptimizers` happens
+  to accept.** `_make_optimizer_cache` and `_make_optimizer_state` asked the capability question
+  (`x isa GeometricOptimizers.OptimizerSolution`, via `_use_go_cache`) *before* the structural one, so
+  a `NetworkParameters` reached the container branch only because it is not currently a member of that
+  union. The moment `GeometricOptimizers` adopts the container — which is the next thing it does — the
+  root of a network would have matched `_use_go_cache` instead, and a whole network would have been
+  given one cache rather than one per layer, with `_leaf_optim_step!` handed the entire tree and
+  `_GMLGradient` handed a `NetworkParameters` it has no method for. A `MethodError` on the first step
+  of every training run, from a change that reads as purely additive upstream.
+
+  The `NetworkParameters` branch now comes first. The `NamedTuple` branch deliberately stays *after*
+  `_use_go_cache`, and that asymmetry is the fix rather than an oversight: a layer is a `NamedTuple` of
+  arrays, which is exactly what one `GeometricOptimizers` cache is for, so hoisting it too would
+  descend into the individual weights. Behaviour today is unchanged, which is what makes it safe to
+  land before the upstream release rather than with it.
+
+  `_tree_optim_step!` had the same inversion in its descent into `λY`, and it is fixed the same way:
+  the test now names both container types, because a tree of sections is something to descend into
+  whichever type carries it. Today only a `NamedTuple` arrives, since this package's own
+  `GlobalSection(::NetworkParameters)` unwraps the container first — but that method is the one it
+  gives back next, and its replacement upstream returns a container, at which point a single
+  `isa NamedTuple` would have handed every layer the whole tree instead of its own section.
+
+  `test/optimizers/utils/optimization_step.jl` is the regression net, and pins the shape rather than
+  the run: a network's cache and state are `NamedTuple`s keyed by its layers, with one
+  `GeometricOptimizers` cache and one state per layer — not one for the root, and not one per weight.
+  `test/optimizers/structured_array_parameters.jl`, four architectures × four methods, covers the step
+  itself.
+
+### Documentation
+
+- `_tree_optim_step!` records why it is *not* written with
+  `NeuralNetworkParameters.foreachparameters`, having been an obvious candidate. It walks the **cache**
+  tree, which stops at the layer where a cache sits, whereas `foreachparameters` walks the leaf
+  protocol and would descend past the layer into individual weights and re-pair every cache with the
+  wrong object. And `λY` is broadcast rather than zipped — a single `GlobalSection` may stand in for a
+  whole subtree — which `foreachparameters` cannot express, because it takes `values` of each trailing
+  argument. The `nothing`-skip is the only thing the two have in common, and it is one line here.
+
 ## [0.6.0] — 2026-08-24
 
 **The parameter container moves out to [NeuralNetworkParameters.jl][nnp].**

README.md

@@ -30,7 +30,7 @@ include("scripts/pendulum.jl")
 
 type = Float32 # Float16 etc.
 # get data 
-qp_data = GeometricMachineLearning.apply_toNT(a -> CuArray(type.(a)), pendulum_data((q=[0.], p=[1.]); timespan=(0.,100.)))
+qp_data = map(a -> CuArray(type.(a)), pendulum_data((q=[0.], p=[1.]); timespan=(0.,100.)))
 # call the DataLoader
 dl = DataLoader(qp_data)
 

src/GeometricMachineLearning.jl

@@ -5,7 +5,13 @@ using AbstractNeuralNetworks
 # under the name `NetworkParameters`. The import is selective rather than a bare `using`: that
 # package also exports `flatten`/`unflatten` and the leaf protocol, none of which this package
 # extends — `GeometricOptimizers` carries the protocol for the structured matrices.
+#
+# Two things do come from there rather than being written again here: `mapstorage`, which reaches the
+# storage of a structured leaf and rebuilds the leaf around the result (`src/map_to_cpu.jl`), and
+# `parameter_eltype`, which promotes over the leaves of a set. A plain `NamedTuple` of layers is
+# walked with `Base.map`, which needs nothing from anybody.
 import NeuralNetworkParameters: NetworkParameters
+using NeuralNetworkParameters: mapstorage, parameter_eltype
 using ChainRulesCore
 # `sqeuclidean` is the default distance of every `TrainingMethod` in `src/training_method/`.
 using Distances
@@ -166,8 +172,10 @@ include("activations/softmax.jl")
 # are these needed?
 export UnknownProblem, NothingFunction
 
-# + operation has been overloaded to work with NamedTuples!
-export _add, apply_toNT, add!
+# `_add`, `_diff` and `_norm` are the `NamedTuple`/`(q, p)` arms of addition, subtraction and the
+# norm, and none of the three is exported: they are helpers of `src/reduced_system/`, not surface.
+# `_add` was the odd one out until 0.7.0, as was `add!` -- which is `AbstractNeuralNetworks`' generic
+# and available from there, this package only adding methods for the structured matrix types.
 
 # GPU specific operations
 export convert_to_dev, Device, CPUDevice

src/map_to_cpu.jl

@@ -1,34 +1,18 @@
-function map_to_cpu(ps::NetworkParameters) 
-    NetworkParameters(NamedTuple{keys(ps)}(Tuple(map_to_cpu(ps[key]) for key in keys(ps))))
-end
-
-map_to_cpu(layer::NamedTuple) = apply_toNT(map_to_cpu, layer)
-
-function map_to_cpu(A::AbstractArray{T}) where T
-    Array{T}(A)
-end
-
-function map_to_cpu(Y::StiefelManifold{T}) where T 
-    StiefelManifold(Array{T}(Y.A))
-end
-
-function map_to_cpu(U::UpperTriangular{T}) where T
-    UpperTriangular(Array{T}(U.S), U.n)
-end
-
-function map_to_cpu(L::LowerTriangular{T}) where T
-    LowerTriangular(Array{T}(L.S), L.n)
-end
-
-function map_to_cpu(A::SkewSymMatrix{T}) where T
-    SkewSymMatrix(Array{T}(A.S), A.n)
-end
-
-function map_to_cpu(A::SymmetricMatrix{T}) where T 
-    SymmetricMatrix(Array{T}(A.S), A.n)
-end
+# Move a set of parameters, or a whole network, from a device back to the host.
+#
+# One walk covers every leaf. `mapstorage` hands `f` the `freeparameters` of a leaf and `rebuild`s the
+# leaf around the result, so a `StiefelManifold` comes back a `StiefelManifold` and a `SymmetricMatrix`
+# keeps its `n` -- which is precisely what the five per-type methods this replaces were doing by hand.
+# `GeometricOptimizers` supplies the protocol for its own structured types, so nothing here has to know
+# which of them exist, and a type added upstream is covered without a change on this side.
+#
+# `mapstorage` and not `mapparameters`: the latter hands `f` *whole* leaves, which would still need one
+# method per structured type to reach the storage.
+_to_host(A::AbstractArray{T}) where {T} = Array{T}(A)
+
+map_to_cpu(ps) = mapstorage(_to_host, ps)
 
 function map_to_cpu(nn::NeuralNetwork{AT, MT, <:Any, BT}) where {AT, MT, BT}
     ps = map_to_cpu(params(nn))
     NeuralNetwork{AT, MT, typeof(ps), BT}(nn.architecture, nn.model, ps, nn.backend)
-end
+end

src/nnsolution/history.jl

@@ -52,7 +52,7 @@ end
 Base.getindex(history::History, n::Int) = hdata(history)[n]
 Base.iterate(history::History, state = 1) = state > size(history) ? nothing : (history[state],state+1)
 
-function _add(history::History, sg::SingleHistory)
+function _push_history!(history::History, sg::SingleHistory)
 
     history.nbtraining += 1 
 

src/nnsolution/neural_net_solution.jl

@@ -25,7 +25,7 @@ struct NeuralNetSolution{TNN <: NeuralNetwork, TP <: AbstractProblem, timestep <
     end
 end
 
-update_history(nns::NeuralNetSolution, sg::SingleHistory) = _add(nns.history, sg)
+update_history(nns::NeuralNetSolution, sg::SingleHistory) = _push_history!(nns.history, sg)
 
 @inline nn(nns::NeuralNetSolution) = nns.nn
 @inline problem(nns::NeuralNetSolution) = nns.problem

src/optimizers/optimizer.jl

@@ -15,8 +15,7 @@ end
 
 _gml_rgrad(x::Manifold, dp) = rgrad(x, dp)
 _gml_rgrad(x, dp) = dp
-_gml_rgrad(x::NamedTuple, dp::NamedTuple) =
-    GeometricOptimizers.apply_toNT(_gml_rgrad, x, dp)
+_gml_rgrad(x::NamedTuple, dp::NamedTuple) = map(_gml_rgrad, x, dp)
 
 (g::_GMLGradient{T})(x::GeometricOptimizers.ArrayNamedTuple{T}) where {T} =
     _gml_rgrad(x, g.dp)
@@ -52,20 +51,39 @@ _adapt_method_to_T(method, ::Type) = method
 _use_go_cache(method, x) =
     _is_go_native_method(method) && x isa GeometricOptimizers.OptimizerSolution
 
+# A `NetworkParameters` is always a tree of layers to descend into, never a single
+# `GeometricOptimizers` leaf, so its branch comes *first* — ahead of `_use_go_cache`.
+#
+# Today that ordering is invisible: `NetworkParameters` is not one of the types
+# `GeometricOptimizers.OptimizerSolution` unions, so `_use_go_cache` is false at the root anyway and
+# control reaches the container branch either way. It stops being invisible the moment
+# `GeometricOptimizers` adopts the container and adds it to that union. Then `_use_go_cache` would be
+# true at the *root*, and a whole network would get one cache instead of one per layer -- silently, and
+# with `_leaf_optim_step!` handed the entire tree. Asking the structural question before the
+# capability question is what makes the shape of the cache depend on the shape of the parameters
+# rather than on which types upstream happens to accept this month.
+#
+# The `NamedTuple` branch stays *after* `_use_go_cache`, and that asymmetry is the point: a layer is a
+# `NamedTuple` of arrays, which is exactly what one `GeometricOptimizers` cache is for. Hoisting it
+# too would descend into the individual weights and give each its own `GMLEuclideanState`.
 function _make_optimizer_cache(method, x)
-    if _use_go_cache(method, x)
-        GeometricOptimizers.OptimizerCache(_adapt_method_to_T(method, _eltype(x)), x)
-    elseif x isa NamedTuple || x isa NetworkParameters
+    if x isa NetworkParameters
+        NamedTuple{keys(x)}(Tuple(_make_optimizer_cache(method, x[k]) for k in keys(x)))
+    elseif _use_go_cache(method, x)
+        GeometricOptimizers.OptimizerCache(_adapt_method_to_T(method, parameter_eltype(x)), x)
+    elseif x isa NamedTuple
         NamedTuple{keys(x)}(Tuple(_make_optimizer_cache(method, x[k]) for k in keys(x)))
     else
         GMLEuclideanState(x)
     end
 end
 
 function _make_optimizer_state(method, x)
-    if _use_go_cache(method, x)
-        GeometricOptimizers.OptimizerState(_adapt_method_to_T(method, _eltype(x)), x)
-    elseif x isa NamedTuple || x isa NetworkParameters
+    if x isa NetworkParameters
+        NamedTuple{keys(x)}(Tuple(_make_optimizer_state(method, x[k]) for k in keys(x)))
+    elseif _use_go_cache(method, x)
+        GeometricOptimizers.OptimizerState(_adapt_method_to_T(method, parameter_eltype(x)), x)
+    elseif x isa NamedTuple
         NamedTuple{keys(x)}(Tuple(_make_optimizer_state(method, x[k]) for k in keys(x)))
     else
         GMLEuclideanState(x)
@@ -197,7 +215,7 @@ end
 
 function _go_update_leaf!(cache, state, local_grad,
         method::GeometricOptimizers.OptimizerMethod, ps_leaf)
-    T = _eltype(ps_leaf)
+    T = parameter_eltype(ps_leaf)
     GeometricOptimizers.update!(cache, state, local_grad,
                                 GeometricOptimizers.NoHessian{T}(), ps_leaf)
 end
@@ -206,7 +224,7 @@ end
 function _leaf_optim_step!(cache::GeometricOptimizers.OptimizerCache,
         state::GeometricOptimizers.OptimizerState,
         dp_leaf, ps_leaf, λY_leaf, method, retraction, step_size)
-    T = _eltype(ps_leaf)
+    T = parameter_eltype(ps_leaf)
     local_grad = _GMLGradient{T, typeof(dp_leaf)}(dp_leaf)
     adapted = _adapt_method_to_T(method, T)
     state.iterations += 1
@@ -248,13 +266,35 @@ function _leaf_optim_step!(cache::GMLEuclideanState, state::GMLEuclideanState,
     nothing
 end
 
-# Recursive dispatcher over the parameter tree
+# Recursive dispatcher over the *cache* tree.
+#
+# Deliberately hand-written rather than `NeuralNetworkParameters.foreachparameters`, which walks a
+# parameter tree. Two reasons, and both are load-bearing:
+#
+#  - The recursion is keyed on `caches`, and stops where the cache stops. A cache sits at the *layer*
+#    level, so a layer's `NamedTuple` of weights arrives at `_leaf_optim_step!` whole, which is what
+#    one `GeometricOptimizers` cache is for. `foreachparameters` recurses on the leaf protocol and
+#    would descend past the layer into the individual weights, re-pairing every cache with the wrong
+#    object.
+#  - `λY` is broadcast, not zipped: a single `GlobalSection` may stand in for a whole subtree, which
+#    the ternary below expresses. `foreachparameters` has no such rule — it takes `values` of each
+#    trailing argument, so a bare `GlobalSection` beside a `NamedTuple` of caches is a `MethodError`.
+#
+# The `nothing` skip is the one thing the two have in common, and it is one line here.
+#
+# The `λY` test names both container types for the same reason `_make_optimizer_cache` asks the
+# structural question first: a section tree is something to descend into, whichever type carries it.
+# Today only a `NamedTuple` ever arrives, because the `GlobalSection(::NetworkParameters)` method at
+# the top of this file unwraps the container before `GeometricOptimizers` sees it. That method is
+# GML's to give back once `GeometricOptimizers` depends on `NeuralNetworkParameters` — and its
+# replacement there is expected to return a *container* of sections, at which point `isa NamedTuple`
+# alone would be false and every layer would be handed the whole tree instead of its own section.
 function _tree_optim_step!(caches, states, dp, ps, λY, method, retraction, step_size)
     if caches isa NamedTuple
         for k in keys(caches)
             dp_k = dp[k]
             dp_k === nothing && continue
-            λY_k = λY isa NamedTuple ? λY[k] : λY
+            λY_k = λY isa Union{NamedTuple, NetworkParameters} ? λY[k] : λY
             _tree_optim_step!(caches[k], states[k], dp_k, ps[k], λY_k,
                               method, retraction, step_size)
         end