Train on GeometricOptimizers 0.9: NetworkOptimizer and train! (part M1a)

CHANGELOG.md

@@ -15,8 +15,32 @@ breaking release).
 
 ## [Unreleased]
 
+### Removed (breaking)
+
+- **`Optimizer`, the callable optimizers and their aliases are gone; training uses GeometricOptimizers' `TrainingOptimizer`.** `TrainingOptimizer`, re-exported by GeometricMachineLearning from GeometricOptimizers 0.9, is built on the parameters of a network: `TrainingOptimizer(nn.params; algorithm = Adam(), linesearch = default_step_size(algorithm), retraction = Cayley())`. Migrate `Optimizer(method, nn; step_size = η, retraction = cayley)` to `TrainingOptimizer(nn.params; algorithm = method, linesearch = η, retraction = Cayley())`, and `Optimizer(nn; AdamOptimizerWithDecay(n, T)...)` to `TrainingOptimizer(nn.params; AdamOptimizerWithDecay(n)...)`. The `n` of `AdamOptimizerWithDecay` counts optimization steps, one per minibatch. One cache and one state now cover the whole parameter set; 0.8 had one per layer. `Optimizer` is no longer exported by GeometricMachineLearning, so with both packages loaded, `Optimizer` is GeometricOptimizers' solver.
+
+- **`train!` replaces the four callable methods of `Optimizer`.** `train!(nn, opt, dl, batch, n_epochs; loss, pullback = ZygotePullback(loss), show_progress = true)` replaces `(o::Optimizer)(nn, dl, batch, n_epochs[, loss, pullback])` and `optimize_for_one_epoch!`, which is removed from the exports. Migrate `o(nn, dl, batch, n, loss; kw...)` to `train!(nn, o, dl, batch, n; loss, kw...)`. `loss` is a `NetworkLoss`, and `pullback` is any `AbstractPullback`: a Hamiltonian network trains with `loss = HNNLoss(arch)` and `pullback = SymbolicPullback(arch)`. The sixth positional `AbstractPullback` that 0.8 allowed is gone. Without `loss`, a `TransformerIntegrator` with a `Batch{:Transformer}` trains on `TransformerLoss(batch)`, a `NeuralNetworkIntegrator` with a `Batch{:FeedForward}` on `FeedForwardLoss()`, and an autoencoder with a `Batch{:FeedForward}` on `AutoEncoderLoss()`; any other combination raises an `ArgumentError` that asks for `loss`. The `λY` argument and the per-call `GlobalSection` are gone: `opt` holds the sections, moments and iteration number, so a second `train!` call with the same `opt` resumes the first, and `train!` draws no random numbers beyond the minibatch shuffle. `n_epochs < 1` raises an `ArgumentError` before any step. The loss history has the element type of the loss value, so a `Float64` network on `Float32` data gives a `Vector{Float64}`.
+
+- **A hand-written loop steps with `GeometricOptimizers.optimization_step!(nn.params, opt, dp)`.** The 0.8 method `optimization_step!(opt, λY, ps, dp)` is removed. GeometricMachineLearning imports `optimization_step!` but does not export it, so a caller qualifies it.
+
+- **A gradient leaf that is all zero is stepped, not skipped.** NeuralNetworkParameters 0.5 gives a layer the loss does not read a zero leaf of its own type, not `nothing`, and the step takes the whole gradient. A leaf the loss never reads stays where it is under `GradientMethod`, `MomentumMethod` and `Adam`, and a Stiefel leaf stays under all four methods; under `AdamWithEuclideanDecay` a Euclidean leaf decays. A leaf the loss reads for some steps and then not keeps moving on its momentum under `MomentumMethod`, `Adam` and `AdamWithEuclideanDecay`, and stays under `GradientMethod`. `scripts/verification/all_zero_leaf.jl` asserts them.
+
+- **This closes known issues C1, C2 and D4.** C1 and C2 close because GML's own parameter-tree traversal and `_leaf_optim_step!` are deleted; D4 closes because the per-layer caches it measured are deleted.
+
+- **The aliases `GradientOptimizer`, `MomentumOptimizer` and `AdamOptimizer` are removed**; use `GradientMethod`, `MomentumMethod` and `Adam`. These are no longer exported: `OptimizerMethod`, `GradientState`, `MomentumState`, `AdamState`, `optimization_step!`, `optimize_for_one_epoch!` and `Optimizer`. `TrainingOptimizer` (re-exported from GeometricOptimizers), `train!` and `AdamWithEuclideanDecay` are newly exported, and the package re-exports `GradientMethod`, `MomentumMethod`, `Adam`, `AdamOptimizerWithDecay`, `DecayingStatic`, `Cayley` and `Geodesic` from GeometricOptimizers. Methods carry no element type in GeometricOptimizers 0.9: `Adam(T)` becomes `Adam()`, and `MomentumMethod(α)` becomes `MomentumMethod(; α)`. `check(::Optimizer)` is removed.
+
+- **The three 3-tensor functors of `Dense` and `Linear` and the HDF5 `save` and `load` methods on `NeuralNetwork` are removed**, because AbstractNeuralNetworks 0.9 provides them. The HDF5 extension and the HDF5 weak dependency are gone; with HDF5 loaded, AbstractNeuralNetworks' extension provides `save` and `load`. Its `load` has no `backend` keyword: load on the CPU, then call `changebackend(backend, nn)`. Aqua's `piracies` check runs again with 0 piracies (was 3), and the ambiguity count is 0 (was 1). This closes issue #145 and known issue B7.
+
+### Changed (breaking)
+
+- **`LowerTriangular` and `UpperTriangular` are renamed `StrictlyLowerTriangular` and `StrictlyUpperTriangular`**, as in GeometricOptimizers 0.9, and are re-exported under the new names. A file saved with the old names no longer loads through the registry. The rigid-body tutorial registers the two old names for its committed files.
+
+- **The `[compat]` bounds move to the versions the training API needs**: GeometricOptimizers `"0.9"`, AbstractNeuralNetworks `"0.9.1"`, NeuralNetworkParameters `"0.5"`, SymbolicNeuralNetworks `"0.9"`, GeometricBase `"0.15"`, GeometricEquations `"0.21.5"` and GeometricSolutions `"0.6.6"`. HDF5 leaves `[compat]` and `[weakdeps]`.
+
 ### Added
 
+- **`KernelAbstractions.get_backend` is defined for `PoissonTensor`.** GeometricOptimizers 0.9 asks every operand of a checked product for its backend, so `𝕁 * Y` with a `StiefelManifold` raised `ArgumentError: Implement KernelAbstractions.get_backend(::PoissonTensor...)` without it.
+
 - **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.
@@ -27,7 +51,7 @@ breaking release).
 - `Pkg.test()` no longer runs the doctests. `test/quality/doctests.jl` is now the group `doctests`, which an empty `ARGS` does not run; `Pkg.test(test_args = ["doctests"])` runs it. In CI the Doctests job stays their runner, so the test matrix no longer runs them a second time.
 - **CI uploads coverage from the `Julia 1 - ubuntu-latest` job instead of `Julia min`, and a test job saves the Julia cache only when it succeeds.**
 - **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.
-- **`test/Project.toml` has no `[compat]` entry for a dependency of the package.** The entries for AbstractNeuralNetworks, GeometricOptimizers, HDF5, KernelAbstractions, LinearAlgebra, NeuralNetworkParameters, Random and Zygote are gone: the test environment contains the package, so the bounds of `Project.toml` apply there, and an entry in `test/Project.toml` can only repeat or narrow them. `NeuralNetworkParameters = "0.3"` was narrower than the package's `"0.3, 0.4"`.
+- **`test/Project.toml` has no `[compat]` entry for a dependency of the package.** The entries for AbstractNeuralNetworks, GeometricOptimizers, KernelAbstractions, LinearAlgebra, NeuralNetworkParameters, Random and Zygote are gone: the test environment contains the package, so the bounds of `Project.toml` apply there, and an entry in `test/Project.toml` can only repeat or narrow them. `NeuralNetworkParameters = "0.3"` was narrower than the package's `"0.3, 0.4"`.
 - **The test suite follows one flat layout.** `test/runtests.jl` is one list of `@safetestset "<label>" include("<path>")` lines in the groups `core`, `slow` and `metal`, selected by `ARGS`; with no arguments it runs `core` and `slow`, and `metal` too on an Apple-silicon Mac, so `Pkg.test()` and CI run every test they ran before. The 14 nested `test/<subject>/runtests.jl` drivers are gone. A test file sits where the source it tests sits: `test/attention/` and `test/transformers/` went into `test/layers/`, `test/architectures/` and `test/kernels/`; `test/reduced_order_modeling/` into `test/architectures/` and `test/reduced_system/`; `test/losses/training_method_losses.jl` into `test/loss/`; the tests that span the package, the docstring examples among them, into `test/integration/`; and `aqua.jl`, `exports.jl` and `reachability.jl` into `test/quality/`. `test/docstrings/layers_and_architectures.jl`, named in *Added* above, is now `test/integration/docstrings/layers_and_architectures.jl`. `test/layers/sympnet_upscaling.jl` is in `slow`. Thirteen files that drew random numbers without a seed now call `Random.seed!(1234)`. The shared helper of the two gradient-structure tests is `test/helpers/network_parameters_gradient_structure.jl`, and its `@test_broken` names issue #319.
 - **The doctests run in the test suite.** `test/quality/doctests.jl`, in the `slow` group, runs `doctest(GeometricMachineLearning)` as the `Doctests` job of `CI.yml` does. `Documenter` joins `test/Project.toml`, with no `[compat]` bound, as a test-only dependency.
 - **The test suite compiles less.** `test/layers/sympnet_upscaling.jl` checks its three symplectic identities on 12 pairs `(N, N2)` instead of all 120 with `N` in `2:2:20` and `N2` in `(2N):2:(4N)`: `(2, 4)`, `(2, 6)`, `(2, 8)`, `(4, 10)`, `(6, 12)`, `(6, 20)`, `(8, 24)`, `(10, 30)`, `(12, 36)`, `(14, 42)`, `(20, 40)` and `(20, 80)`. They keep the smallest pair, `N2 = 2N`, `N2 = 4N`, `N2` not a multiple of `N`, and the largest `N`. The layers carry the sizes in their types, so each pair is a new compilation (about 1.3 s per pair, not once per file as the [0.8.0] entry assumed); the file has 36 assertions instead of 360, and takes 19.7 s cold; it stays in `slow`. The two training tests of `test/loss/training_method_losses.jl` use the network shape of the rest of the file, `StandardHamiltonianArchitecture(2n, 4, 1)` and `LagrangianNeuralNetwork(2; width = 4, nhidden = 1)`, in place of `(2n, 5, 2)` and `width = 5, nhidden = 2`; each still asserts that training reduces the loss. `test/architectures/lagrangian_neural_network_tests.jl` moves from `core` to `slow`: it takes 122.7 s cold, above the 60 s of a `core` file. It still runs in CI and in `Pkg.test()`, and leaves only the local `core` run. Every other `core` file takes at most 37.9 s cold. The times are cold, from one fresh Julia process that runs the `core` files in the order of `test/runtests.jl` (and from one fresh process for each of the two files named), on an Apple M4 Max with Julia 1.13.1, at a load average below 4 at the start.