Adopt AbstractNeuralNetworks 0.7: NetworkParameters, and the traversal with it

CHANGELOG.md

@@ -36,6 +36,39 @@ breaking release).
   of the module already provides, so `GeometricMachineLearning.IdentityActivation` still resolves —
   the line was redundant, not load-bearing.
 
+- **`NeuralNetworkParameters` is no longer exported; the name is `NetworkParameters`.** The parameter
+  container moved out of `AbstractNeuralNetworks` into
+  [NeuralNetworkParameters.jl](https://github.com/JuliaGNI/NeuralNetworkParameters.jl) in
+  `AbstractNeuralNetworks` 0.7, which removed the old name outright rather than leaving an alias, so
+  that one type has one name across the ecosystem. This package follows, at every call site — 31 in
+  `src/`, `docs/` and `test/`, seven more in `scripts/` — and in its export list. It is the same type
+  object, so `::Type{}` dispatch, `<:` bounds and `NetworkParameters{keys}(vals)` construction are
+  unaffected; only the spelling changes.
+
+  ```julia
+  # before
+  using GeometricMachineLearning        # brought NeuralNetworkParameters into scope
+  # after
+  using GeometricMachineLearning        # brings NetworkParameters into scope
+  ```
+
+- **The HDF5 extension no longer carries its own traversal.** Five `h5save` methods tagging a
+  `gml_type` attribute, the `_gml_h5load` reader and the `_natural_sort_keys` key-order heuristic are
+  gone — 73 of the extension's 187 lines. Each job now sits with the package that owns the pieces:
+
+  - `NeuralNetworkParameters` walks the parameter set and writes it, recording each group's key order
+    in a `keys` attribute. `_natural_sort_keys` was standing in for that, and it *guessed*: it sorted
+    on a trailing integer when every name in the group had one and fell back to lexicographic order
+    otherwise, so a group whose names do not end in a digit came back in whatever order sorting gave.
+  - `GeometricOptimizers` says where each structured matrix keeps its numbers, through
+    `freeparameters`/`rebuild`, and registers the types so a file loads with no prototype.
+    `StiefelManifold` and `SymmetricMatrix` are its types, not this package's, so the methods here
+    were type piracy twice over — on `h5save` and on the type.
+
+  Existing files still load. `NeuralNetworkParameters` recognises the `gml_type` tag and rebuilds
+  through the same registry, and `test/hdf5_support.jl` now writes a file in the old layout by hand
+  and reads it back, so the deletion cannot quietly make old files unreadable.
+
 ### Fixed
 
 - **Zygote 0.7 silently zeroed every gradient that flows through `assign_q_and_p`.** Its `rrule`
@@ -102,8 +135,39 @@ breaking release).
   it sits in, on a signature it did not need — the body ignores its argument and returns
   `ZeroTangent()` regardless.
 
+- **`save(filename, nn)` returns `filename`.** It used to return whatever the `h5open` block left
+  behind — the value of the innermost `h5save`, an implementation detail of the traversal. Returning
+  the path is what `NeuralNetworkParameters.save(filename, ps)` does, so the two now agree.
+
+### Added
+
+- **`load(NeuralNetwork, h5, arch, prototype)`** — a parameter set of the right shape to rebuild the
+  structured leaves against. It is the form that needs no registration: `rebuild` has a prototype to
+  take the non-differentiable fields from, so the file's type tags and
+  `NeuralNetworkParameters.register_parameter_type!` are not consulted at all. Both the store and the
+  filename overloads take it.
+
 ### Dependencies
 
+- **`NeuralNetworkParameters = "0.1"`** added, and **`AbstractNeuralNetworks = "0.7"`** (was
+  `"0.6.4"`). The parameter container is defined in the former as of the latter; see *Removed* above.
+
+- **`GeometricOptimizers = "0.4.1"`** (was `"0.4"`). 0.4.1 is the release that carries the
+  `NeuralNetworkParameters` leaf protocol for the manifolds, storage matrices and horizontal lifts,
+  which is what lets this package's HDF5 extension drop its own copy of the traversal.
+
+- **`SymbolicNeuralNetworks = "0.6"`** (was `"0.5"`). 0.5 caps `AbstractNeuralNetworks` at `"0.6.4 -
+  0.6"`, so leaving the bound would have made this package's `[compat]` unsatisfiable rather than
+  merely unresolved. 0.6 is the release that follows the container out to `NeuralNetworkParameters`.
+
+  > **Merge order.** `AbstractNeuralNetworks` 0.7.0, `GeometricOptimizers` 0.4.1 and
+  > `NeuralNetworkParameters` 0.1.1 are all in the General registry as of 2026-08-23.
+  > `SymbolicNeuralNetworks` 0.6.0 is not: its `abstractneuralnetworks-0.7` branch still says
+  > `0.5.0`, and still does `using AbstractNeuralNetworks: QPTOAT`, which 0.7 replaced with
+  > `ArrayOrNamedTuple`, so it does not load as it stands. That is the one release this waits on.
+  >
+  > Until it lands, CI here fails at `Pkg.instantiate`. That is expected, not a regression.
+
 - **`Zygote = "0.7"`** (was `"0.6"`). 0.7 replaced the eager unthunking in `wrap_chainrules_output`
   with `unthunk_tangent` at the `gradient`/`pullback` boundaries, which is what let thunks reach
   GML's `rrule`s and surfaced everything under *Fixed* above. Implicit parameters are deprecated in

Project.toml

@@ -1,6 +1,6 @@
 name = "GeometricMachineLearning"
 uuid = "194d25b2-d3f5-49f0-af24-c124f4aa80cc"
-version = "0.5.0"
+version = "0.6.0"
 authors = ["Michael Kraus <michael.kraus@ipp.mpg.de>"]
 
 [deps]
@@ -16,6 +16,7 @@ InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+NeuralNetworkParameters = "67f4d93a-60e9-472b-8cdd-1ccf6005724a"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SymbolicNeuralNetworks = "aed23131-dcd0-47ca-8090-d21e605652e3"
@@ -30,22 +31,23 @@ HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 HDF5Ext = "HDF5"
 
 [compat]
-AbstractNeuralNetworks = "0.6.4"
+AbstractNeuralNetworks = "0.7"
 ChainRulesCore = "1"
 ChainRulesTestUtils = "1"
 Distances = "0.10"
 ForwardDiff = "1"
 GeometricBase = "0.14"
 GeometricEquations = "0.21"
 GeometricIntegrators = "0.18.2"
-GeometricOptimizers = "0.4"
+GeometricOptimizers = "0.4.1"
 GeometricSolutions = "0.6"
 HDF5 = "0.16, 0.17"
 KernelAbstractions = "0.9"
 NNlib = "0.8, 0.9"
+NeuralNetworkParameters = "0.1"
 ProgressMeter = "1"
 SafeTestsets = "0.1"
-SymbolicNeuralNetworks = "0.5"
+SymbolicNeuralNetworks = "0.6"
 Symbolics = "7"
 TimerOutputs = "0.5, 1"
 Zygote = "0.7"

docs/src/architectures/abstract_neural_networks.md

@@ -47,7 +47,7 @@ and we see that it consists of two layers: a [`GradientLayerQ`](@ref) and a [`Gr
 
 ## Saving and Loading
 
-`GeometricMachineLearning` extends `AbstractNeuralNetworks` with HDF5-backed save and load methods for `NeuralNetwork`. These handle GML-specific parameter types (`StiefelManifold`, `SymmetricMatrix`, `SkewSymMatrix`) transparently.
+`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)

docs/src/tutorials/adjusting_the_loss_function.md

@@ -42,7 +42,7 @@ using LinearAlgebra: norm  # hide
 # norm of parameters for single layer
 network_parameter_norm(params::NamedTuple) = sum([norm(params[i]) for i in 1:length(params)])
 # norm of parameters for entire network
-function network_parameter_norm(params::NeuralNetworkParameters)
+function network_parameter_norm(params::NetworkParameters)
     sum([network_parameter_norm(params[key]) for key in keys(params)])
 end
 
@@ -60,7 +60,7 @@ struct CustomLoss <: GeometricMachineLearning.NetworkLoss end
 using GeometricMachineLearning: QPTOAT, AbstractExplicitLayer # hide
 
 const λ = .1
-function (loss::CustomLoss)(model::Union{AbstractExplicitLayer, Chain}, params::Union{NeuralNetworkParameters, NamedTuple}, input::QPTOAT, output::QPTOAT)
+function (loss::CustomLoss)(model::Union{AbstractExplicitLayer, Chain}, params::Union{NetworkParameters, NamedTuple}, input::QPTOAT, output::QPTOAT)
     FeedForwardLoss()(model, params, input, output) + λ * network_parameter_norm(params)
 end
 nothing # hide

docs/src/tutorials/grassmann_layer.md

@@ -260,7 +260,7 @@ where `np` is the number of points in ``\mathcal{D}_2`` and ``W_2`` is the *Wass
 where ``\nabla{}W_2`` is equivalent to the function `compute_wasserstein_gradient`.
 
 ```@example rosenbrock
-function compute_gradient(ps::NeuralNetworkParameters)
+function compute_gradient(ps::NetworkParameters)
     samples = randn(2, size(xyz_points, 2))
     estimate, nn_pullback = Zygote.pullback(ps -> model(samples, ps), ps)
 

ext/HDF5Ext.jl

@@ -2,58 +2,22 @@ module HDF5Ext
 
 using HDF5
 using GeometricMachineLearning
-import AbstractNeuralNetworks: h5save, changebackend, NeuralNetworkBackend, save, load,
-                                NeuralNetworkParameters, params, Architecture
-
-# ---------------------------------------------------------------------------
-# h5save — new methods for GML special array types
-#
-# AbstractNeuralNetworks only defines h5save for AbstractArray and NamedTuple.
-# PSDLayer stores StiefelManifold and LASympNet's LinearLayer stores
-# SymmetricMatrix; without these methods h5save throws a MethodError on any
-# network whose parameters include those types.
-# ---------------------------------------------------------------------------
-
-function h5save(h5::HDF5.H5DataStore, Y::StiefelManifold, path::AbstractString)
-    group = haskey(h5, path) ? h5[path] : HDF5.create_group(h5, path)
-    HDF5.attributes(group)["gml_type"] = "StiefelManifold"
-    group["A"] = Array(Y.A)
-end
-
-function h5save(h5::HDF5.H5DataStore, A::SymmetricMatrix, path::AbstractString)
-    group = haskey(h5, path) ? h5[path] : HDF5.create_group(h5, path)
-    HDF5.attributes(group)["gml_type"] = "SymmetricMatrix"
-    group["S"] = Array(A.S)
-    group["n"] = A.n
-end
-
-function h5save(h5::HDF5.H5DataStore, A::SkewSymMatrix, path::AbstractString)
-    group = haskey(h5, path) ? h5[path] : HDF5.create_group(h5, path)
-    HDF5.attributes(group)["gml_type"] = "SkewSymMatrix"
-    group["S"] = Array(A.S)
-    group["n"] = A.n
-end
-
-function h5save(h5::HDF5.H5DataStore, A::LowerTriangular, path::AbstractString)
-    group = haskey(h5, path) ? h5[path] : HDF5.create_group(h5, path)
-    HDF5.attributes(group)["gml_type"] = "LowerTriangular"
-    group["S"] = Array(A.S)
-    group["n"] = A.n
-end
-
-function h5save(h5::HDF5.H5DataStore, A::UpperTriangular, path::AbstractString)
-    group = haskey(h5, path) ? h5[path] : HDF5.create_group(h5, path)
-    HDF5.attributes(group)["gml_type"] = "UpperTriangular"
-    group["S"] = Array(A.S)
-    group["n"] = A.n
-end
+import AbstractNeuralNetworks: changebackend, NeuralNetworkBackend, Architecture
+# `save`, `load`, `params` and the parameter container are `NeuralNetworkParameters`' as of
+# `AbstractNeuralNetworks` 0.7, which only re-binds them; reach for them where they are defined.
+import NeuralNetworkParameters: NetworkParameters, params, save, load
 
 # ---------------------------------------------------------------------------
 # changebackend — new methods for GML special array types
 #
 # AbstractNeuralNetworks.changebackend handles AbstractArray and NamedTuple.
 # Moving a NeuralNetwork between devices fails for parameters that include
 # StiefelManifold, SymmetricMatrix, or SkewSymMatrix without these methods.
+#
+# `changebackend` is `AbstractNeuralNetworks`' and the types are `GeometricOptimizers`', so these
+# methods are piracy the same way the `h5save` ones were before `GeometricOptimizers` took over the
+# leaf protocol. They belong in a `GeometricOptimizers` extension on `AbstractNeuralNetworks`; that
+# is a separate change with its own release chain, so they stay here for now.
 # ---------------------------------------------------------------------------
 
 function changebackend(backend::NeuralNetworkBackend, Y::StiefelManifold)
@@ -77,105 +41,66 @@ function changebackend(backend::NeuralNetworkBackend, A::UpperTriangular)
 end
 
 # ---------------------------------------------------------------------------
-# Internal recursive loader that reconstructs GML special types from the
-# gml_type attribute written by h5save.  Kept private; used only by the
-# load methods below so that we do not shadow AbstractNeuralNetworks.h5load.
-# ---------------------------------------------------------------------------
-
-_gml_h5load(ds::HDF5.Dataset) = read(ds)
-
-# HDF5 returns group keys alphabetically, so "L10" precedes "L2". Sort by the
-# numeric suffix when all keys match the pattern <letters><digits>, otherwise
-# fall back to lexicographic order so per-layer NamedTuples (bias/scale/weight)
-# are unaffected.
-function _natural_sort_keys(ks)
-    if all(k -> occursin(r"^\D+\d+$", k), ks)
-        return sort(collect(ks),
-            by = k -> (m = match(r"^(\D+)(\d+)$", k); (m[1], parse(Int, m[2]))))
-    end
-    sort(collect(ks))
-end
-
-function _gml_h5load(group::HDF5.Group)
-    if haskey(HDF5.attributes(group), "gml_type")
-        gml_type = read(HDF5.attributes(group)["gml_type"])
-        if gml_type == "StiefelManifold"
-            return StiefelManifold(read(group["A"]))
-        elseif gml_type == "SymmetricMatrix"
-            return SymmetricMatrix(read(group["S"]), read(group["n"]))
-        elseif gml_type == "SkewSymMatrix"
-            return SkewSymMatrix(read(group["S"]), read(group["n"]))
-        elseif gml_type == "LowerTriangular"
-            return LowerTriangular(read(group["S"]), read(group["n"]))
-        elseif gml_type == "UpperTriangular"
-            return UpperTriangular(read(group["S"]), read(group["n"]))
-        end
-    end
-    sorted_keys = _natural_sort_keys(keys(group))
-    paramkeys = Tuple(Symbol.(sorted_keys))
-    paramvals = Tuple(_gml_h5load(group[k]) for k in sorted_keys)
-    NamedTuple{paramkeys}(paramvals)
-end
-
-# ---------------------------------------------------------------------------
-# save — new dispatch on NeuralNetwork, mirroring the existing
-#   save(h5::H5DataStore, p::NeuralNetworkParameters)
-# method in AbstractNeuralNetworks.
+# save / load — the entry points that dispatch on this package's `NeuralNetwork`.
+#
+# The traversal itself is not here. `NeuralNetworkParameters` walks the parameter set and records
+# each group's key order; `GeometricOptimizers` says through `freeparameters`/`rebuild` where each
+# structured matrix keeps its numbers, and registers the types so a file loads with no prototype.
 # ---------------------------------------------------------------------------
 
 """
     save(h5::HDF5.H5DataStore, nn::NeuralNetwork)
 
 Save the parameters of `nn` into an already-open HDF5 store.
 
-Extends `AbstractNeuralNetworks.save` with a dispatch on `NeuralNetwork`.
-GML special array types (`StiefelManifold`, `SymmetricMatrix`, `SkewSymMatrix`,
-`LowerTriangular`, `UpperTriangular`) are tagged with a `gml_type` attribute
-so that [`load`](@ref) can reconstruct them faithfully.
+Extends `save` with a dispatch on `NeuralNetwork`. The parameters themselves are written by
+`NeuralNetworkParameters`, which tags each structured leaf with the type to rebuild it as and
+records the key order of every group.
 """
-function save(h5::HDF5.H5DataStore, nn::NeuralNetwork)
-    h5save(h5, params(params(nn)), "/")
-end
+save(h5::HDF5.H5DataStore, nn::NeuralNetwork) = save(h5, params(nn))
 
 """
     save(filename::AbstractString, nn::NeuralNetwork)
 
-Convenience overload: open `filename` for writing, then call
-`save(h5, nn)`.
+Convenience overload: open `filename` for writing, call [`save`](@ref) on the store, and return
+`filename`.
 """
 function save(filename::AbstractString, nn::NeuralNetwork)
     HDF5.h5open(filename, "w") do h5
         save(h5, nn)
     end
+    filename
 end
 
-# ---------------------------------------------------------------------------
-# load — new dispatch on NeuralNetwork, mirroring the existing
-#   load(::Type{NeuralNetworkParameters}, h5::H5DataStore)
-# method in AbstractNeuralNetworks.
-# ---------------------------------------------------------------------------
-
 """
-    load(::Type{NeuralNetwork}, h5::HDF5.H5DataStore, arch::Architecture; backend=CPU())
+    load(::Type{NeuralNetwork}, h5::HDF5.H5DataStore, arch::Architecture; backend = CPU())
+    load(::Type{NeuralNetwork}, h5::HDF5.H5DataStore, arch::Architecture, prototype; backend = CPU())
 
-Load network parameters from an already-open HDF5 store and return a
-`NeuralNetwork` for `arch`.
+Load network parameters from an already-open HDF5 store and return a `NeuralNetwork` for `arch`.
 
-Extends `AbstractNeuralNetworks.load` with a dispatch on `NeuralNetwork`.
-GML special array types are reconstructed from their `gml_type` attribute.
-The element type is preserved as stored (Float32 files reload as Float32).
+The element type is whatever the file holds, so a `Float32` network reloads as `Float32`.
+
+Structured parameters — `StiefelManifold`, `SymmetricMatrix` and the rest — are rebuilt from the
+type each was stored under, which `GeometricOptimizers` registers with
+`NeuralNetworkParameters.register_parameter_type!`. Pass `prototype`, a parameter set of the right
+shape, to rebuild against it instead and skip the registry altogether.
 """
 function load(::Type{NeuralNetwork}, h5::HDF5.H5DataStore, arch::Architecture;
               backend::NeuralNetworkBackend = CPU())
-    ps = NeuralNetworkParameters(_gml_h5load(h5["/"]))
-    NeuralNetwork(arch, Chain(arch), ps, backend)
+    NeuralNetwork(arch, Chain(arch), load(NetworkParameters, h5), backend)
+end
+
+function load(::Type{NeuralNetwork}, h5::HDF5.H5DataStore, arch::Architecture, prototype;
+              backend::NeuralNetworkBackend = CPU())
+    NeuralNetwork(arch, Chain(arch), load(NetworkParameters, h5, prototype), backend)
 end
 
 """
-    load(::Type{NeuralNetwork}, filename::AbstractString, arch::Architecture; backend=CPU())
+    load(::Type{NeuralNetwork}, filename::AbstractString, arch::Architecture; backend = CPU())
+    load(::Type{NeuralNetwork}, filename::AbstractString, arch::Architecture, prototype; backend = CPU())
 
 Convenience overload: open `filename` for reading, then call
-`load(NeuralNetwork, h5, arch; backend)`.
+[`load`](@ref) on the store.
 """
 function load(::Type{NeuralNetwork}, filename::AbstractString, arch::Architecture;
               backend::NeuralNetworkBackend = CPU())
@@ -184,4 +109,11 @@ function load(::Type{NeuralNetwork}, filename::AbstractString, arch::Architectur
     end
 end
 
+function load(::Type{NeuralNetwork}, filename::AbstractString, arch::Architecture, prototype;
+              backend::NeuralNetworkBackend = CPU())
+    HDF5.h5open(filename, "r") do h5
+        load(NeuralNetwork, h5, arch, prototype; backend = backend)
+    end
+end
+
 end