Wrap a layer at the GeometricOptimizers boundary, and drop ParameterSet

CHANGELOG.md

@@ -13,45 +13,12 @@ breaking release).
 > alongside the work. Where a release removed exported names the list is given; where it is a
 > reconstruction of intent, it says so.
 
-## [Unreleased] — 0.6.1
+## [Unreleased] — 0.8.0
 
-### Changed
-
-- **Julia 1.11 is the minimum**, up from the 1.10 LTS, in step with the rest of this family of
-  packages. Nothing here was accommodating 1.10; the two places that name it —
-  `test/training_parameters.jl` and `test/optimizers/optimizer_convergence_tests/psd_optim.jl` — record
-  where a measurement was taken and how a defect was found, and both stay as the history they are.
-
-- **`GeometricOptimizers = "0.6"`.** 0.6.0 is what takes a whole `NetworkParameters` through the
-  optimizer, and it is breaking: four `outer!`/`_mul!` methods, one `update!(::BFGSState, …)` and
-  `add!` on a parameter set are gone, and two `l2norm` methods moved upstream to `GeometricBase`. None
-  of them had a caller here — the `add!` arm was deleted on this side in 0.7 for the same reason — so
-  the bump is a compat entry and no code. The full suite is green against it.
-
-- **The fifteen sites spelling `Union{NamedTuple, NetworkParameters}` inline now name it.** It is
-  `NeuralNetworkParameters.ParameterSet`, added upstream in 0.2.2: the same union, in the package that
-  owns the type. `SymbolicNeuralNetworks` had it as `EquationSet`, `AbstractNeuralNetworks` spelled it
-  out eight times and `GeometricOptimizers` sixteen; all four name one thing now. Definitionally
-  identical, so no dispatch changes anywhere. Compat is `NeuralNetworkParameters = "0.2.2"`.
-
-### Fixed
-
-- **The comment on `_make_optimizer_cache`'s branch ordering had been wrong for a release.** It said
-  `NetworkParameters` is not one of the types `GeometricOptimizers.OptimizerSolution` unions, and
-  therefore that asking the structural question before the capability question was invisible. 0.5.0 made
-  the container a member, so `_use_go_cache` is true at the *root* now and the ordering is load-bearing:
-  without it a whole network would get one cache instead of one per layer, silently, with
-  `_leaf_optim_step!` handed the entire tree. The code was already right — the ordering was put in
-  ahead of the change it anticipated — and only the comment needed correcting. Found in the review of
-  [GeometricOptimizers #68](https://github.com/JuliaGNI/GeometricOptimizers.jl/pull/68).
-
-  It also records what the better end state is, now that `GeometricOptimizers` 0.6.0 takes a whole
-  container as one solution: one cache for the network, which means giving `_GMLGradient` a
-  `NetworkParameters` method and dropping `_tree_optim_step!`. That is a change of behaviour — one
-  `GlobalSection` tree and one `Q` across every layer instead of one per layer — so it wants its own
-  release rather than being folded in here.
-
-## [0.7.0]
+> [!NOTE]
+> Not released. This was written as `[0.7.0]` before 0.6.1 and 0.7.0 were cut, and the number was
+> taken by the release below; the work it describes -- `apply_toNT`, `_eltype`, `map_to_cpu` -- is
+> still to do.
 
 **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
@@ -215,6 +182,135 @@ so it cannot coexist with `GeometricOptimizers` 0.5.
   docs build — here, one release behind. The old one-liner also called `save` unqualified, and
   `DocInventories` does not export it.
 
+
+## [0.7.0]
+
+**A layer is wrapped at the `GeometricOptimizers` boundary, and a whole set of parameters is a
+`NetworkParameters`.** Breaking: this release tracks `NeuralNetworkParameters` 0.3.0,
+`AbstractNeuralNetworks` 0.8.0, `SymbolicNeuralNetworks` 0.8.0 and `GeometricOptimizers` 0.7.0.
+
+### Changed
+
+- **`GeometricOptimizers = "0.7"`.** 0.6.0 is what takes a whole `NetworkParameters` through the
+  optimizer, and it was breaking: four `outer!`/`_mul!` methods, one `update!(::BFGSState, …)` and
+  `add!` on a parameter set are gone, and two `l2norm` methods moved upstream to `GeometricBase`. None
+  of them had a caller here — the `add!` arm was deleted on this side in 0.7 for the same reason.
+
+  0.7.0 goes further and **does** need code here: a whole set of parameters reaches that package only
+  as a `NetworkParameters` now, and a bare `NamedTuple` is turned away. This package hands it one
+  **layer** at a time, so it cannot convert by wrapping at the root; see the entry below. The full
+  suite is green against it, `test/reduced_system.jl` aside — that one needs `GeometricIntegrators`,
+  which still pins `SimpleSolvers` 0.12 and is the ecosystem-wide blocker.
+
+- **`ParameterSet` is gone from the signatures here.** `NeuralNetworkParameters` 0.3.0 removes it, so
+  `Optimizer` and `optimize_for_one_epoch!` take a `NetworkParameters`, `_get_contents` has a method
+  per shape, and the loss functors drop the annotation entirely — it dispatched nothing there, since
+  the model and the input/output types settle every method, and both a container and the bare
+  `NamedTuple` a reverse pass produces reach them.
+
+  `_tree_optim_step!`'s test for a section tree is now `_is_section_tree`, a named predicate rather than
+  an `isa` against a union. It answers one question — descend, or hand this whole section to the layer
+  — and naming it is what keeps the container arm from looking like breadth for its own sake.
+
+- **A *flat* set of parameters gets one optimizer cache again, and the branch order is what decides it.**
+  `_make_optimizer_cache` asks `_is_layer` — "is this one cache's worth?" — *before* asking whether the
+  argument is a tree to descend into. A flat set answers yes to both, and the tree reading is wrong: it
+  gives every individual weight its own cache, and for a manifold weight that is not merely wasteful,
+  because a bare `Manifold` then reaches `_GMLGradient` where a whole layer should have, which is
+  ambiguous against `GeometricOptimizers`' own `Gradient` functor. `_is_layer` accepts a container as
+  well as a bare `NamedTuple` for this reason.
+
+- **A layer is wrapped at the `GeometricOptimizers` boundary, and the flatness rule is written out
+  here.** `_as_go_solution` wraps a layer in a `NetworkParameters` before handing it to
+  `OptimizerCache`, `OptimizerState` or `update!`, and `_leaf_optim_step!` uses the wrapped form
+  throughout. The wrap **shares the leaf arrays**, so the step still writes through to the network's
+  own weights and nothing is copied back — which is what makes this a boundary change and not a
+  change of behaviour.
+
+  What did have to be restated is the test that separates a *layer* from a *subtree*. Upstream made it
+  by dispatch: `ArrayNamedTuple{T}` bounded a `NamedTuple`'s values by `AbstractArray{T}`, so a
+  `NamedTuple` of branches was not one and got a cache per layer rather than one for the whole of it.
+  That alias is gone — it was an alias for `Base.NamedTuple`, so every method on it was a method on a
+  `Base` type — so `_is_layer` says it here, where it is this package's rule rather than a property of
+  whichever types upstream happens to accept this month.
+
+  **One behaviour follows from the wrap rather than from the old bound.** A layer whose weights do not
+  share an element type is now one cache with a promoted `T`, where the `Vararg` bound admitted no `T`
+  at all and the layer fell through to one `GMLEuclideanState` per weight — i.e. no manifold handling.
+  `_adapt_method_to_T` reads the promotion, so the cache is built for the type the layer actually has.
+
+  `_GMLGradient` gains its `NetworkParameters` method as part of this, and `_gml_rgrad` walks with
+  `mapparameters` instead of `Base.map`: it recurses on the branches, so `_gml_rgrad` is only ever
+  called on leaves, and it rebuilds in the shape of its first argument — a container, which is what
+  `GeometricOptimizers._copyto!(gradient_array(cache), ·)` has a method for.
+
+### Fixed
+
+- **`_GMLGradient` on a bare `Manifold` was ambiguous, and the branch ordering was the only thing
+  preventing it.** `_GMLGradient{T}` is a `GeometricOptimizers.Gradient{T}` and `Manifold{T}` is an
+  `AbstractMatrix{T}`, so `(::_GMLGradient{T})(::AbstractArray{T})` here and upstream's
+  `(::Gradient{T})(::Manifold{T})` are neither of them more specific on the intersection —
+  `MethodError: … is ambiguous`, several frames into a solve. Asking `_is_layer` before descending
+  keeps a bare `Manifold` from arriving, and that ordering is a correctness fix in its own right, but
+  a routing decision must not be the only thing standing between a caller and an ambiguity.
+
+  There is now a method for it, and the answer it gives is this package's: the Euclidean gradient is
+  already computed and carried in `dp`, so the projection is `rgrad(x, dp)` rather than upstream's
+  `rgrad(x, reshape(grad(vec(x)), …))`, which would evaluate an inner gradient this functor does not
+  have. `test/optimizers/gml_gradient_dispatch.jl` pins it, along with the shapes the optimizer
+  actually produces.
+
+- **An empty set was one cache's worth.** `all` over an empty collection is vacuously true, so
+  `_is_layer(NamedTuple())` was `true` and `OptimizerCache` would have been asked for the element type
+  of a set with no leaves to promote. An empty set is *zero* caches' worth and descends.
+
+- `_get_contents` takes a method per shape in its wrapped forms too, rather than a `Tuple` and an
+  `AbstractVector` over a union of the two. Same rule as everywhere else in this release.
+
+## [0.6.1]
+
+### Changed
+
+- **Julia 1.11 is the minimum**, up from the 1.10 LTS, in step with the rest of this family of
+  packages. Nothing here was accommodating 1.10; the two places that name it —
+  `test/training_parameters.jl` and `test/optimizers/optimizer_convergence_tests/psd_optim.jl` — record
+  where a measurement was taken and how a defect was found, and both stay as the history they are.
+
+- **`GeometricOptimizers = "0.6"`.** 0.6.0 is what takes a whole `NetworkParameters` through the
+  optimizer, and it is breaking: four `outer!`/`_mul!` methods, one `update!(::BFGSState, …)` and
+  `add!` on a parameter set are gone, and two `l2norm` methods moved upstream to `GeometricBase`. None
+  of them had a caller here — the `add!` arm was deleted on this side in 0.7 for the same reason — so
+  the bump is a compat entry and no code. The full suite is green against it.
+
+- **The fifteen sites spelling `Union{NamedTuple, NetworkParameters}` inline now name it.** It is
+  `NeuralNetworkParameters.ParameterSet`, added upstream in 0.2.2: the same union, in the package that
+  owns the type. `SymbolicNeuralNetworks` had it as `EquationSet`, `AbstractNeuralNetworks` spelled it
+  out eight times and `GeometricOptimizers` sixteen; all four name one thing now. Definitionally
+  identical, so no dispatch changes anywhere. Compat is `NeuralNetworkParameters = "0.2.2"`.
+- **The fifteen sites spelling `Union{NamedTuple, NetworkParameters}` inline now name it.** It is
+  `NeuralNetworkParameters.ParameterSet`, added upstream in 0.2.2: the same union, in the package that
+  owns the type. `SymbolicNeuralNetworks` had it as `EquationSet`, `AbstractNeuralNetworks` spelled it
+  out eight times and `GeometricOptimizers` sixteen; all four name one thing now. Definitionally
+  identical, so no dispatch changes anywhere. Compat is `NeuralNetworkParameters = "0.2.2"`.
+
+### Fixed
+
+- **The comment on `_make_optimizer_cache`'s branch ordering had been wrong for a release.** It said
+  `NetworkParameters` is not one of the types `GeometricOptimizers.OptimizerSolution` unions, and
+  therefore that asking the structural question before the capability question was invisible. 0.5.0 made
+  the container a member, so `_use_go_cache` is true at the *root* now and the ordering is load-bearing:
+  without it a whole network would get one cache instead of one per layer, silently, with
+  `_leaf_optim_step!` handed the entire tree. The code was already right — the ordering was put in
+  ahead of the change it anticipated — and only the comment needed correcting. Found in the review of
+  [GeometricOptimizers #68](https://github.com/JuliaGNI/GeometricOptimizers.jl/pull/68).
+
+  It also records what the better end state is, now that `GeometricOptimizers` takes a whole container
+  as one solution: one cache for the network. `_GMLGradient` has its `NetworkParameters` method as of
+  this release, but that is one wrapped **layer** and not the tree — the container branch still splits
+  the network first. Dropping `_tree_optim_step!` and handing the whole tree to one cache is a change
+  of behaviour — one `GlobalSection` tree and one `Q` across every layer instead of one per layer — so
+  it still wants its own release rather than being folded in here.
+
 ## [0.6.0] — 2026-08-24
 
 **The parameter container moves out to [NeuralNetworkParameters.jl][nnp].**

Project.toml

@@ -1,6 +1,6 @@
 name = "GeometricMachineLearning"
 uuid = "194d25b2-d3f5-49f0-af24-c124f4aa80cc"
-version = "0.6.1"
+version = "0.7.0"
 authors = ["Michael Kraus <michael.kraus@ipp.mpg.de>"]
 
 [deps]
@@ -31,23 +31,23 @@ HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 HDF5Ext = "HDF5"
 
 [compat]
-AbstractNeuralNetworks = "0.7"
+AbstractNeuralNetworks = "0.8"
 ChainRulesCore = "1"
 ChainRulesTestUtils = "1"
 Distances = "0.10"
 ForwardDiff = "1"
 GeometricBase = "0.14"
 GeometricEquations = "0.21"
 GeometricIntegrators = "0.18.2"
-GeometricOptimizers = "0.6"
+GeometricOptimizers = "0.7"
 GeometricSolutions = "0.6"
 HDF5 = "0.16, 0.17"
 KernelAbstractions = "0.9"
 NNlib = "0.8, 0.9"
-NeuralNetworkParameters = "0.2.2"
+NeuralNetworkParameters = "0.3"
 ProgressMeter = "1"
 SafeTestsets = "0.1"
-SymbolicNeuralNetworks = "0.7"
+SymbolicNeuralNetworks = "0.8"
 Symbolics = "7"
 TimerOutputs = "0.5, 1"
 Zygote = "0.7"

src/GeometricMachineLearning.jl

@@ -11,7 +11,7 @@ using AbstractNeuralNetworks
 # `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, ParameterSet
+using NeuralNetworkParameters: mapparameters, mapstorage, parameter_eltype
 using ChainRulesCore
 # `sqeuclidean` is the default distance of every `TrainingMethod` in `src/training_method/`.
 using Distances

src/data_loader/optimize.jl

@@ -48,7 +48,7 @@ number_of_batches(dl, batch)
 ```
 """
 function optimize_for_one_epoch!(   opt::Optimizer, 
-                                    model, ps::ParameterSet, 
+                                    model, ps::NetworkParameters, 
                                     dl::DataLoader{T}, 
                                     batch::Batch, 
                                     loss::NetworkLoss, 
@@ -58,7 +58,7 @@ end
 
 function optimize_for_one_epoch!(   opt::Optimizer, 
                                     model, 
-                                    ps::ParameterSet, 
+                                    ps::NetworkParameters, 
                                     dl::DataLoader{T}, 
                                     batch::Batch, 
                                     _pullback::AbstractPullback, 

src/loss/hnn_loss.jl

@@ -29,13 +29,13 @@ end
 AbstractNeuralNetworks.NetworkLoss(arch::HamiltonianArchitecture) = HNNLoss(arch)
 
 function (loss::HNNLoss)(::Union{Chain, AbstractExplicitLayer},
-        ps::ParameterSet,
+        ps,
         input::QPTOAT,
         output::QPTOAT)
     loss(ps, input, output)
 end
 
-function (loss::HNNLoss)(ps::ParameterSet,
+function (loss::HNNLoss)(ps,
         input::QPTOAT,
         output::QPTOAT)
     norm(loss.hvf(input, ps) - output) / norm(output)

src/loss/losses.jl

@@ -73,7 +73,7 @@ function crop_array_for_transformer_loss(nn_output::AT,
 end
 
 function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer},
-        ps::ParameterSet, input::AT,
+        ps, input::AT,
         output::AT) where {T, AT <: AbstractArray{T, 3}}
     input_dim, input_seq_length = size(input)
     output_dim, output_prediction_window = size(output)
@@ -88,13 +88,13 @@ function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer},
 end
 
 function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer},
-        ps::ParameterSet, input::AT,
+        ps, input::AT,
         output::AT) where {T, AT <: AbstractArray{T, 2}}
     loss(model, ps, reshape(input, size(input)..., 1), reshape(output, size(output)..., 1))
 end
 
 function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer},
-        ps::ParameterSet,
+        ps,
         input::T, output::T) where {T <: QPT}
     loss(model, ps, vcat(input.q, input.p), vcat(output.q, output.p))
 end
@@ -120,7 +120,7 @@ end
 struct ClassificationTransformerLoss <: NetworkLoss end
 
 function (loss::ClassificationTransformerLoss)(model::Union{Chain, AbstractExplicitLayer},
-        ps::ParameterSet,
+        ps,
         input::AbstractArray, output::AbstractArray)
     predicted_output_uncropped = model(input, ps)
     # predicted_output_cropped = crop_array_for_transformer_loss(predicted_output_uncropped, output)
@@ -179,12 +179,12 @@ function (loss::AutoEncoderLoss)(nn::NeuralNetwork, input::QPTOAT)
 end
 
 function (loss::AutoEncoderLoss)(model::Union{Chain, AbstractExplicitLayer},
-        ps::ParameterSet, input::QPTOAT)
+        ps, input::QPTOAT)
     loss(model, ps, input, input)
 end
 
 function (loss::AutoEncoderLoss)(model::Union{Chain, AbstractExplicitLayer},
-        ps::ParameterSet, input::QPTOAT, output::QPTOAT)
+        ps, input::QPTOAT, output::QPTOAT)
     FeedForwardLoss()(model, ps, input, output)
 end