Train through GeometricOptimizers' TrainingOptimizer (G4), one method per leaf

CHANGELOG.md

@@ -17,11 +17,75 @@ breaking release).
 
 ### Added
 
+- **One optimizer method per leaf, through `CompositeMethod`.** `CompositeMethod` and
+  `CompositeState` are re-exported from `GeometricOptimizers`, with `ScalarMomentAdam` and
+  `ScalarMomentAdamState`. `Optimizer(CompositeMethod(; manifold = ScalarMomentAdam(), array = Adam()), nn)`
+  gives every leaf of the network's parameters the cache and the state of the method chosen for it,
+  which is what a network with Stiefel weights beside Euclidean ones needs to use `ScalarMomentAdam`:
+  that method steps a single `StiefelManifold`. A caller that wanted the split before had to add
+  methods to `_make_optimizer_cache`, `_make_optimizer_state` and `_leaf_optim_step!` from outside
+  this package. `test/optimizers/composite_method.jl` checks the per-leaf states and that a composite
+  of `Adam` on every leaf reproduces `Adam` exactly.
 - **The test suite runs on an Apple GPU in its `metal` group.** `test/devices/metal.jl` checks the two device paths that were measured by hand and not tested before: the products of a `PoissonTensor` with a wrapped device array, `𝕁 * view(A, [1, 2, 3, 4], :)` and `𝕁 * B'`, which `ext/GPUArraysCoreExt.jl` carries; and the `tensor_mat_mul!` kernel on `MetalBackend()`. Each result must stay on the device and match the host in `Float32`, with scalar indexing off. `Pkg.test(test_args = ["metal"])` runs them without the subject tests, and a missing device records a skip, which happens inside a sandbox. Metal (1.10 or later) is a test dependency on every platform; it installs and precompiles on Linux and Windows, and nothing there loads it. On an Apple M4 Max with Metal 1.11.1 all five assertions pass.
 - **A `Metal` workflow runs the Metal tests on GitHub's `macos-15` runner**, for the `min` and `1` Julia versions. The runner's GPU is Apple's paravirtual device, which Metal.jl supports from 1.10 and on macOS 15 or later. The workflow includes a step "Metal is functional" before the tests that fails the job where the device is unavailable, so it cannot pass without having run the tests. The job is not a required check. However, the Metal tests also run in two required checks—CI.yml's `Julia min - macOS-latest - default` and `Julia 1 - macOS-latest - default`—which execute `test/runtests.jl` with no test arguments on Apple silicon, selecting the `metal` group.
 - **`test/quality/jet.jl` checks the kernel launchers with JET.** It runs in `core`, directly after `quality/aqua.jl`, and takes about 14 s. Each function that launches a kernel of `src/` on the `CPU()` backend, the pullbacks among them, and the `Batch` functor that `test/data_loader/batch_index_set.jl` measures with `@allocated` get one `JET.report_opt` line per element type that the tests use, with `target_modules = (GeometricMachineLearning,)`. Each `@kernel` method also gets a line that analyses its generated `cpu_<kernel>` function at the launcher's context, because JET drops the reports of a kernel body when it analyses the launcher. 142 lines pass on Julia 1.11.9 (JET 0.9.20) and 1.13.1 (JET 0.12.2). Three lines report a runtime dispatch on both versions and are `@test_broken` with issue #325: `PoissonTensor(::CPU, ::Int, ::DataType)`, and the input-and-output method of `convert_input_and_batch_indices_to_array` at `Float32` and `Float64`. Where JET does not work, the file records one `@test_skip`. JET joins `test/Project.toml` with no `[compat]` bound. `KNOWN_ISSUES.md` K2 records the Revise `EMFILE` messages that JET brings into the test log, and K3 the dispatches that the lines do not see: an unstable argument of a kernel launch, and two barriers in the Poisson-tensor kernel.
 - **A test calls `iterate` on `(q,)` initial conditions.** No test reached the `NeuralNetworkIntegrator` `iterate` method that takes `ics::NamedTuple{(:q,), Tuple{AT}}`. `test/docstrings/layers_and_architectures.jl` runs the docstring's `ResNet` from `(q = [1, 1, 1],)` for four points and expects the same trajectory as the vector form, `[1 2 4 8; 1 3 9 27; 1 3 7 15]`.
 
+### Changed (breaking)
+
+- **The training step is `GeometricOptimizers.TrainingOptimizer`.** `Optimizer` holds one, over the
+  whole parameter set of the network, in its only field `training`; the fields `method`, `cache`,
+  `state`, `retraction`, `step_size` and `iterations` are gone. One cache and one state cover the
+  whole set, where there used to be one per layer, and every method steps every kind of leaf the same
+  way `GeometricOptimizers` does: no leaf is stepped by this package's own Euclidean update rules
+  (`GMLEuclideanState`) any more. `GeometricOptimizers`' G4 checked that `TrainingOptimizer`
+  reproduces this package's 0.8 iterates: exactly for `Adam`, and to `4eps(T)` for `GradientMethod` and
+  `MomentumMethod` on manifold weights. `ScalarMomentAdam` trains the Stiefel weights of a network
+  through a `CompositeMethod` (see *Added*); the step used to raise a `MethodError`, since this
+  package had no method of its own for it.
+- **`optimization_step!(ps, opt, dp)` replaces `optimization_step!(opt, λY, ps, dp)`**, and
+  `optimize_for_one_epoch!` loses its last argument `λY`. The section is carried by the optimizer's
+  state, so the caller no longer builds a `GlobalSection` of the parameters for it.
+  `optimization_step!` is `GeometricOptimizers`' function, extended for `Optimizer`, so a `Main` that
+  loads both packages has one `optimization_step!` and not two. `dp` is the gradient a pullback
+  returns, a `NetworkParameters` or the bare `NamedTuple` of the same shape.
+- **`retraction` takes `Cayley()` or `Geodesic()`**, the retraction types, and no longer the
+  functions `cayley` and `geodesic`. The default is `Cayley()`.
+- **The methods carry no element type**, as in `GeometricOptimizers` 0.9: `Adam()`,
+  `Adam(; β₁, β₂, δ)`, `MomentumMethod(; α)` and `AdamOptimizerWithDecay(n_epochs; η₁, η₂)` in place
+  of `Adam(Float32)`, `Adam(T; …)`, `MomentumMethod(α)` and `AdamOptimizerWithDecay(n_epochs, T)`.
+  `Optimizer` converts the method to the element type of the parameters.
+- **`LowerTriangular` and `UpperTriangular` are `StrictlyLowerTriangular` and
+  `StrictlyUpperTriangular`**, the names `GeometricOptimizers` 0.9 gives them, which no longer collide
+  with `LinearAlgebra`'s. A file saved with one of them by 0.8 or earlier is tagged with the old name
+  and still loads (JuliaGNI/GeometricOptimizers.jl#146); `test/integration/hdf5_support.jl` loads a
+  `VolumePreservingFeedForward` network written in the old layout.
+- **Requires `GeometricOptimizers` 0.9, `AbstractNeuralNetworks` 0.9, `NeuralNetworkParameters`
+  0.4.2, `GeometricBase` 0.15 and `SymbolicNeuralNetworks` 0.9.** No registered `GeometricOptimizers`
+  before 0.9 allows `NeuralNetworkParameters` 0.4 or `AbstractNeuralNetworks` 0.9, and no registered
+  `SymbolicNeuralNetworks` before 0.9 allows either (JuliaGNI/SymbolicNeuralNetworks.jl#80), so the
+  five move together. Under `NeuralNetworkParameters` 0.4 a Zygote gradient taken through a
+  `NetworkParameters` is the gradient with respect to each leaf's storage: a `SymmetricMatrix` leaf
+  keeps its type however often the loss reads it, which closes #319, and its off-diagonal entries are
+  twice those of the cotangent of the dense matrix. `test/helpers/network_parameters_gradient_structure.jl`
+  asserts both.
+
+### Removed (breaking)
+
+- **The three-axis methods of `Dense` and `Linear`, and the HDF5 extension.** `AbstractNeuralNetworks`
+  0.9 defines both itself: `Dense`, `Linear` and `Affine` take inputs of rank 3 and more, and its
+  `HDF5Ext` has `save` and `load` for a `NeuralNetwork`, so `using HDF5` still loads them. Here they
+  were type piracy, the three methods in `src/layers/resnet.jl` and the four in `ext/HDF5Ext.jl`.
+  Aqua's `piracies` and `ambiguities` checks now run in `test/quality/aqua.jl` and report none; they
+  were switched off with counts of 3 and 1 in their place (*B7* in `KNOWN_ISSUES.md`, now closed).
+  `load(NeuralNetwork, …)` loses its `backend` keyword: the network loads on the CPU, and
+  `changebackend(backend, nn)` moves it.
+- **`GMLEuclideanState`, `_GMLGradient` and the per-layer optimizer walk** (`_make_optimizer_cache`,
+  `_make_optimizer_state`, `_tree_optim_step!`, `_leaf_optim_step!`, `_as_go_solution`) are gone with
+  it, and with them `test/optimizers/gml_gradient_dispatch.jl` and
+  `test/optimizers/step_size_element_type.jl`, which tested two of them. *C2* in `KNOWN_ISSUES.md`
+  is closed; *C1* is narrowed to the one thing left, the name `Optimizer`, which both packages export.
+
 ### Changed
 
 - **The package needs Julia 1.12 or later** (`[compat] julia = "1.12"`, was `"1.11"`). On Julia 1.11.9, the gradient of an `LNNLoss` with respect to the network parameters stops with `StackOverflowError` in type inference ("Internal error: during type inference of … Encountered stack overflow") on macOS and Windows; the `Julia min` job on Linux passes. The overflow comes from SymbolicUtils 4.49.0, whose `rrule` for `SymbolicUtils.Code.create_array` returns an exactly typed tuple of tangents, one per array element; Zygote's reverse pass through the generated Hessian of `lagrangian_acceleration` is then inferred far deeper. With the 4.48.0 rule restored, or in a task with a 64 MB stack, the same gradient passes on 1.11.9; Julia 1.12.7 and 1.13.1 pass it on the root task. This closes #327: the `Julia min` jobs on macOS and Windows had failed in the *Lagrangian Neural Network* testset since 2026-10-01, and every testset after it was never run there.

KNOWN_ISSUES.md

@@ -45,37 +45,6 @@ What is known to be wrong in GeometricMachineLearning and is not fixed.
   every type parameter, and Aqua's `unbound_args` passes on Julia nightly. See *Fixed* above. The
   number is left vacant rather than reused.)
 
-### B7 · Three methods are type piracy, and closing them is an API change rather than a deletion.
-
-- **location:** `src/layers/resnet.jl:63`
-- **kind:** defect
-- **found:** 2026-09-19
-- **evidence:**
-  Aqua's `piracies` check reports 3, down from the 12 this entry opened with; the other nine were
-  deleted in this release and are under *Removed (breaking)* above. All three are genuine under
-  Aqua's definition — the function and every argument type belong to other modules — and all three
-  are one thing: a functor on a *layer* type this package does not own, applied to a three-axis
-  array.
-
-  | method | `src` | without GML | with GML |
-  |:--|:--|:--|:--|
-  | `(::Dense{M, N, true})(::AbstractArray{T, 3}, ::NamedTuple)` | `layers/resnet.jl:63` | `MethodError` | the layer applied along the third axis |
-  | `(::Dense{M, N, false})(::AbstractArray{T, 3}, ::NamedTuple)` | `layers/resnet.jl:67` | the same | the same |
-  | `(::Linear{M, N})(::AbstractArray{T, 3}, ::NamedTuple)` | `layers/resnet.jl:71` | the same | the same |
-
-  `Dense` and `Linear` are `AbstractNeuralNetworks`', and upstream a 3-tensor argument is a
-  `MethodError` because `*` cannot take one. So the witness is not a changed answer but a new one:
-  a call that throws without this package loaded returns a value with it.
-
-  Closing this means either `AbstractNeuralNetworks` gaining the three-axis methods, or this
-  package wrapping the two layer types in its own. Both are an API change, and neither is a
-  by-product of a piracy pass — which is why these three stayed when the other nine went.
-
-  `src/layers/resnet.jl:63` is also the one remaining method ambiguity, against
-  `AbstractNeuralNetworks.Affine`. That pair is benign and has no witness at all: `Dense` is not a
-  subtype of `Affine` and the two have no common instance, so no call can reach it. `test/quality/aqua.jl`
-  asserts both counts — 3 piracies and 1 ambiguity — so neither can drift, in either direction.
-
 ### B10 · `DataLoader(::EnsembleSolution{T, T1, Vector{ST}})` at `src/data_loader/data_loader.jl:325` has no test and appears unreachable through this package's current dependencies.
 
 - **location:** `src/data_loader/data_loader.jl:325`
@@ -135,36 +104,18 @@ What is known to be wrong in GeometricMachineLearning and is not fixed.
 
 ## C. Follow-ups and cleanups
 
-### C1 · The parameter-tree traversal still belongs upstream.
+### C1 · `Optimizer` is exported by both this package and GeometricOptimizers.
 
-- **location:** —
+- **location:** `src/optimizers/optimizer.jl`
 - **kind:** upstream
 - **found:** 2026-08-17
 - **evidence:**
-  `_make_optimizer_cache`,
-  `_make_optimizer_state`, `_tree_optim_step!`, `_leaf_optim_step!` and the bespoke
-  `GMLEuclideanState` are GML implementations of what GeometricOptimizers supports natively for a
-  single parameter. `GMLEuclideanState` in particular duplicates what `GradientState`,
-  `MomentumState` and `AdamState` already do for a plain array.
-
-  What has to go upstream is *not* a reuse of GeometricOptimizers' `Optimizer`: that one needs an
-  `OptimizerProblem`, i.e. an objective function, and minibatch training has none — the gradient
-  arrives from AD one batch at a time. It is a new entry point there, a
-  gradient-supplied-externally step over a `NamedTuple` parameter tree. GML's `Optimizer` would then
-  be the `NeuralNetwork` constructor and the training-loop functor, and nothing else.
-
-  `Optimizer` is the one name still exported by both packages, so this is also what closes the last
-  of B1's class of collision.
-
-### C2 · Two `isa` branches remain in `_leaf_optim_step!`
-
-- **location:** —
-- **kind:** defect
-- **found:** 2026-08-16
-- **evidence:**
-  (for `AdamState`/`MomentumState`).
-  Measurement showed the traversal is not implicated in the compile-time problem, so this is tidying,
-  and it disappears entirely if C1 lands first.
+  The parameter-tree traversal this entry was opened for went upstream in 0.9: the training step is
+  `GeometricOptimizers.TrainingOptimizer`, and GML's `Optimizer` only pairs one with a network and
+  carries the training-loop functor. What remains is the name. GML's `Optimizer` and
+  GeometricOptimizers' `Optimizer` (the one `solve!` takes) are different types, both exported, so a
+  `Main` that `using`s both packages cannot write `Optimizer` unqualified. It is the last of B1's class
+  of collision. Closing it means renaming one of them, which is an API change in either package.
 
 ### C6 · Three generated MNIST PDFs are in this branch's history
 
@@ -363,7 +314,7 @@ Not defects — claims this release makes that nothing has actually checked yet.
 
 ## Found late
 
-### K1 · Four source comments still send the reader to the Open Issues section of `CHANGELOG.md`, which no longer exists.
+### K1 · Two source comments still send the reader to the Open Issues section of `CHANGELOG.md`, which no longer exists.
 
 - **location:** `src/reduced_system/reduced_system.jl:109`
 - **kind:** found late
@@ -373,17 +324,11 @@ Not defects — claims this release makes that nothing has actually checked yet.
   $ grep -rnE 'Open Issues|issue C1[48]|\*C21\*|\*B7\*' src test
   src/kernels/exponentials/tensor_exponential.jl:2:# name is wider than what it holds; renaming it, and `exponentials/` with it, is issue C18.
   src/reduced_system/reduced_system.jl:109:# issue C14 in `CHANGELOG.md`.
-  test/quality/aqua.jl:30:# layer types, and both are an API change rather than a tidy-up. They are *B7* under
-  test/quality/aqua.jl:31:# `## Open Issues` in `CHANGELOG.md` with that reasoning.
   test/reduced_system/reduced_system.jl:44:# for four of twelve, with four `NaN` and four in the opposite order. That is *C21* in
-  $ grep -n 'CHANGELOG' test/quality/aqua.jl
-  31:# `## Open Issues` in `CHANGELOG.md` with that reasoning.
-  63:    # and it fails when one is removed without the entry above and in `CHANGELOG.md` going with it.
   ```
 
-  The IDs are still valid; the comments at `src/reduced_system/reduced_system.jl:109`,
-  `test/quality/aqua.jl:31`, `test/quality/aqua.jl:63` and `test/reduced_system/reduced_system.jl:44–45`
-  name `CHANGELOG.md` rather than `KNOWN_ISSUES.md`. The fix is a follow-up change to those
+  The IDs are still valid; the comments at `src/reduced_system/reduced_system.jl:109` and
+  `test/reduced_system/reduced_system.jl:44–45` name `CHANGELOG.md` rather than `KNOWN_ISSUES.md`. The fix is a follow-up change to those
   comments.
 
 ## Upstream

Project.toml

@@ -25,30 +25,27 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [weakdeps]
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
-HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 
 [extensions]
 GPUArraysCoreExt = "GPUArraysCore"
-HDF5Ext = "HDF5"
 
 [compat]
-AbstractNeuralNetworks = "0.8"
+AbstractNeuralNetworks = "0.9"
 ChainRulesCore = "1"
 ForwardDiff = "1"
 GPUArraysCore = "0.2"
-GeometricBase = "0.14"
+GeometricBase = "0.15"
 GeometricEquations = "0.21"
-GeometricOptimizers = "0.8"
+GeometricOptimizers = "0.9"
 GeometricSolutions = "0.6"
-HDF5 = "0.16, 0.17"
 InteractiveUtils = "1"
 KernelAbstractions = "0.9"
 LinearAlgebra = "1"
 NNlib = "0.8, 0.9"
-NeuralNetworkParameters = "0.3, 0.4"
+NeuralNetworkParameters = "0.4.2"
 ProgressMeter = "1"
 Random = "1"
-SymbolicNeuralNetworks = "0.8"
+SymbolicNeuralNetworks = "0.9"
 Symbolics = "7"
 TimerOutputs = "0.5, 1"
 Zygote = "0.7"

docs/check_references.jl

@@ -12,14 +12,12 @@
 #     julia --project=docs docs/check_references.jl
 
 using GeometricMachineLearning
-using HDF5 # so that the extension below is loaded, as in make.jl
 using AbstractNeuralNetworks # signatures in `@docs` blocks name types from here
 using Documenter
 using Documenter: DocSystem
 
 # keep in sync with `modules` in make.jl
-const MODULES = Module[GeometricMachineLearning,
-    Base.get_extension(GeometricMachineLearning, :HDF5Ext)]
+const MODULES = Module[GeometricMachineLearning]
 
 const SRC = joinpath(@__DIR__, "src")
 

docs/make.jl

@@ -4,6 +4,7 @@
 include(joinpath(@__DIR__, "check_references.jl"))
 
 using GeometricMachineLearning
+# loads `AbstractNeuralNetworks`' HDF5 extension, whose `save` and `load` the tutorials call
 using HDF5
 using AbstractNeuralNetworks
 using Documenter
@@ -387,8 +388,7 @@ makedocs(;
     # skips *sub*modules — and (ii) pull GeometricOptimizers' own docstrings in, whose internal
     # `@ref`s resolve in its namespace and point at bindings this manual does not document.
     # Names owned by GeometricOptimizers are referred to as plain code with a link to its manual.
-    modules = [
-        GeometricMachineLearning, Base.get_extension(GeometricMachineLearning, :HDF5Ext)],
+    modules = [GeometricMachineLearning],
     authors = "Michael Kraus, Benedikt Brantner",
     repo = "https://github.com/JuliaGNI/GeometricMachineLearning.jl/blob/{commit}{path}#L{line}",
     sitename = "GeometricMachineLearning.jl",

docs/src/architectures/abstract_neural_networks.md

@@ -47,11 +47,4 @@ and we see that it consists of two layers: a [`GradientLayerQ`](@ref) and a [`Gr
 
 ## Saving and Loading
 
-`GeometricMachineLearning` adds HDF5-backed `save` and `load` methods for `NeuralNetwork` to the generics `NeuralNetworkParameters` defines. Writing and reading the parameter set itself belongs to that package; the structured parameter types (`StiefelManifold`, `SymmetricMatrix`, `SkewSymMatrix`, …) come back as themselves because `GeometricOptimizers`, which owns them, registers how each is rebuilt. Passing a prototype parameter set to `load` rebuilds against it and needs no registration at all.
-
-```@docs
-save(::HDF5.H5DataStore, ::NeuralNetwork)
-save(::AbstractString, ::NeuralNetwork)
-load(::Type{NeuralNetwork}, ::HDF5.H5DataStore, ::AbstractNeuralNetworks.Architecture)
-load(::Type{NeuralNetwork}, ::AbstractString, ::AbstractNeuralNetworks.Architecture)
-```
+`save` and `load` for a `NeuralNetwork` are `AbstractNeuralNetworks`', which defines them in its HDF5 extension on the generics `NeuralNetworkParameters` provides: `using HDF5` loads them. Writing and reading the parameter set itself belongs to `NeuralNetworkParameters`; the structured parameter types (`StiefelManifold`, `SymmetricMatrix`, `SkewSymMatrix`, …) come back as themselves because `GeometricOptimizers`, which owns them, registers how each is rebuilt. Passing a prototype parameter set to `load` rebuilds against it and needs no registration at all. A network loads on the CPU; `changebackend(backend, nn)` moves it to a device.