From d4e85f3f66a149ba86231a297dd762e3f0465f87 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Mon, 24 Aug 2026 11:56:41 +0900 Subject: [PATCH 1/3] Take the remaining parameter walks from the packages that own them MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `map_to_cpu` loses six per-type methods to `NeuralNetworkParameters.mapstorage`, which reaches the storage of a leaf and rebuilds the leaf around the result. `GeometricOptimizers` supplies that protocol for its own structured types, so nothing here has to know which of them exist. It was untested, so `test/map_to_cpu_tests.jl` comes with it. `apply_toNT` turns out to be `Base.map`: same behaviour at any arity, and Base already rejects mismatched or reordered keys with an `ArgumentError` where the hand-rolled `@assert` could be compiled out. `_norm`, `_diff`, `_add` and `add!`'s container arm call `map` directly, which is the faithful translation rather than a simplification -- two of them recurse through their own `NamedTuple` methods and `_norm` divides by `√length` one level down. `GeometricOptimizers` carries a character-identical copy, reached from here by qualified call; dropping that call is what lets it go in its own release. `_eltype` becomes `parameter_eltype`, which is not a rename: the old one returned the element type of the *first* leaf and read a structured leaf's dense interface, so a layer mixing `Float32` and `Float64` weights picked whichever came first and handed it to `Adam(T)`. Finally, `_make_optimizer_cache` and `_make_optimizer_state` ask the structural question before the capability one. They tested `x isa GeometricOptimizers.OptimizerSolution` first, so a `NetworkParameters` reached the container branch only by virtue of not being in that union yet. Once `GeometricOptimizers` adopts the container, the root of a network would have matched instead and been given one cache rather than one per layer, with `_GMLGradient` handed a container it has no method for. Behaviour today is unchanged, which is what makes it safe to land ahead of that release rather than with it. The `NamedTuple` branch stays after the capability test, and that asymmetry is deliberate: a layer is a `NamedTuple` of arrays, which is exactly what one `GeometricOptimizers` cache is for. --- src/GeometricMachineLearning.jl | 8 +- src/map_to_cpu.jl | 44 ++++------ src/optimizers/optimizer.jl | 54 +++++++++--- src/utils.jl | 24 ++---- test/map_to_cpu_tests.jl | 82 +++++++++++++++++++ test/runtests.jl | 3 + ...ulti_head_attention_stiefel_optim_cache.jl | 6 +- ...multi_head_attention_stiefel_retraction.jl | 4 +- .../multi_head_attention_stiefel_setup.jl | 4 +- 9 files changed, 163 insertions(+), 66 deletions(-) create mode 100644 test/map_to_cpu_tests.jl diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index 5311be0c9..e319bcd7e 100644 --- a/src/GeometricMachineLearning.jl +++ b/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 @@ -167,7 +173,7 @@ include("activations/softmax.jl") export UnknownProblem, NothingFunction # + operation has been overloaded to work with NamedTuples! -export _add, apply_toNT, add! +export _add, add! # GPU specific operations export convert_to_dev, Device, CPUDevice diff --git a/src/map_to_cpu.jl b/src/map_to_cpu.jl index fc5a1a1c3..3173c8226 100644 --- a/src/map_to_cpu.jl +++ b/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 six 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 \ No newline at end of file +end diff --git a/src/optimizers/optimizer.jl b/src/optimizers/optimizer.jl index 5dbba5f17..c41f5619e 100644 --- a/src/optimizers/optimizer.jl +++ b/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,10 +51,27 @@ _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) @@ -63,9 +79,11 @@ function _make_optimizer_cache(method, x) 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,7 +266,21 @@ 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. function _tree_optim_step!(caches, states, dp, ps, λY, method, retraction, step_size) if caches isa NamedTuple for k in keys(caches) diff --git a/src/utils.jl b/src/utils.jl index d4d28a872..bd8456e29 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -28,19 +28,12 @@ develop(t::NamedTuple) = vcat([[develop(e)...] for e in t]...) _tuplediff(t₁::Tuple, t₂::Tuple) = tuple(setdiff(Set(t₁), Set(t₂))...) -function apply_toNT(fun, ps::NamedTuple...) - for p in ps - @assert keys(ps[1]) == keys(p) - end - NamedTuple{keys(ps[1])}(fun(p...) for p in zip(ps...)) -end - # overload norm function _norm(dx::NT) where { AT <: AbstractArray, NT <: NamedTuple{(:q, :p), Tuple{AT, AT}}} (norm(dx.q) + norm(dx.p)) / √2 end # we need this because of a Zygote problem -_norm(dx::NamedTuple) = sum(apply_toNT(norm, dx)) / √length(dx) +_norm(dx::NamedTuple) = sum(map(norm, dx)) / √length(dx) _norm(A::AbstractArray) = norm(A) # overloaded +/- operation @@ -48,9 +41,9 @@ function _diff(dx₁::NT, dx₂::NT) where {AT <: AbstractArray, NT <: NamedTuple{(:q, :p), Tuple{AT, AT}}} (q = dx₁.q - dx₂.q, p = dx₁.p - dx₂.p) end # we need this because of a Zygote problem -_diff(dx₁::NamedTuple, dx₂::NamedTuple) = apply_toNT(_diff, dx₁, dx₂) +_diff(dx₁::NamedTuple, dx₂::NamedTuple) = map(_diff, dx₁, dx₂) _diff(A::AbstractArray, B::AbstractArray) = A - B -_add(dx₁::NamedTuple, dx₂::NamedTuple) = apply_toNT(_add, dx₁, dx₂) +_add(dx₁::NamedTuple, dx₂::NamedTuple) = map(_add, dx₁, dx₂) _add(A::AbstractArray, B::AbstractArray) = A + B function add!(C::AbstractVecOrMat, A::AbstractVecOrMat, B::AbstractVecOrMat) @@ -58,9 +51,10 @@ function add!(C::AbstractVecOrMat, A::AbstractVecOrMat, B::AbstractVecOrMat) C .= A + B end -function add!(dx₁::NamedTuple, dx₂::NamedTuple, dx₃::NamedTuple) - apply_toNT(add!, dx₁, dx₂, dx₃) -end +# No caller in the package; kept because `add!` is `AbstractNeuralNetworks`' generic and this is the +# container arm of it, next to the `AbstractVecOrMat` base case above and the structured-type methods +# in `src/arrays/gml_extensions.jl`. +add!(dx₁::NamedTuple, dx₂::NamedTuple, dx₃::NamedTuple) = map(add!, dx₁, dx₂, dx₃) # Type pyracy!! function Base.:+(a::Float64, b::Tuple{Float64}) @@ -167,7 +161,3 @@ const QPTOAT{T} = Union{QPT{T}, AbstractArray{T}} where {T} Base.:≈(qp₁::QPT, qp₂::QPT) = (qp₁.q ≈ qp₂.q) & (qp₁.p ≈ qp₂.p) -_eltype(x) = eltype(x) -_eltype(ps::NamedTuple) = _eltype(ps[1]) -_eltype(ps::Tuple) = _eltype(ps[1]) -_eltype(ps::NetworkParameters) = _eltype(params(ps)[1]) diff --git a/test/map_to_cpu_tests.jl b/test/map_to_cpu_tests.jl new file mode 100644 index 000000000..677d02536 --- /dev/null +++ b/test/map_to_cpu_tests.jl @@ -0,0 +1,82 @@ +# `map_to_cpu` walks a parameter set with `NeuralNetworkParameters.mapstorage`, so the structured +# types have to survive the round trip rather than come back densified. There is no GPU in CI, so +# what is pinned here is the walk and the reconstruction, not the device transfer: `Array{T}` of a +# host array is a copy, which is enough to show that every leaf was visited and rebuilt. + +using GeometricMachineLearning +using GeometricMachineLearning: map_to_cpu +using LinearAlgebra +using NeuralNetworkParameters: NetworkParameters, params +using Random +using Test + +import GeometricOptimizers: StiefelManifold, SymmetricMatrix, SkewSymMatrix, + LowerTriangular, UpperTriangular + +Random.seed!(1234) + +const N, n = 6, 3 + +@testset "a structured leaf keeps its type and its metadata" begin + ps = NetworkParameters(( + L1 = (Y = rand(StiefelManifold{Float64}, N, n), b = randn(N)), + L2 = (S = SymmetricMatrix(randn(n, n)), A = SkewSymMatrix(randn(n, n))), + L3 = (L = LowerTriangular(randn(n, n)), U = UpperTriangular(randn(n, n))), + )) + + back = map_to_cpu(ps) + + @test back isa NetworkParameters + @test keys(back) == keys(ps) + + # every leaf comes back as the type it went in as -- this is what `rebuild` buys, and what the + # per-type methods this replaced were for + @test back.L1.Y isa StiefelManifold + @test back.L2.S isa SymmetricMatrix + @test back.L2.A isa SkewSymMatrix + @test back.L3.L isa LowerTriangular + @test back.L3.U isa UpperTriangular + + # the `n` a structured leaf carries is not in its storage; `rebuild` takes it from the prototype + @test back.L2.S.n == ps.L2.S.n + @test back.L3.U.n == ps.L3.U.n + + # and the numbers are unchanged + @test back.L1.Y ≈ ps.L1.Y + @test back.L1.b == ps.L1.b + @test back.L2.S == ps.L2.S + @test back.L3.L == ps.L3.L +end + +@testset "the leaves are copies, not the same arrays" begin + ps = NetworkParameters((L1 = (W = randn(2, 2),),)) + back = map_to_cpu(ps) + @test back.L1.W == ps.L1.W + @test back.L1.W !== ps.L1.W +end + +@testset "element type is preserved" begin + ps = NetworkParameters((L1 = (W = randn(Float32, 2, 2), S = SymmetricMatrix(randn(Float32, n, n))),)) + back = map_to_cpu(ps) + @test eltype(back.L1.W) === Float32 + @test eltype(back.L1.S) === Float32 +end + +@testset "a bare NamedTuple of layers works too" begin + ps = (L1 = (W = randn(2, 2),), L2 = (Y = rand(StiefelManifold{Float64}, N, n),)) + back = map_to_cpu(ps) + @test back isa NamedTuple + @test back.L2.Y isa StiefelManifold +end + +@testset "a whole network keeps its architecture, model and backend" begin + nn = NeuralNetwork(Chain(StiefelLayer(N, n), Dense(N, N, tanh))) + back = map_to_cpu(nn) + + @test back.architecture === nn.architecture + @test back.model === nn.model + @test back.backend === nn.backend + @test params(back) isa NetworkParameters + @test keys(params(back)) == keys(params(nn)) + @test params(back).L1.weight isa StiefelManifold +end diff --git a/test/runtests.jl b/test/runtests.jl index d8758cd8d..0adbeae3c 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -101,6 +101,9 @@ end @safetestset "Optimizers with structured (non-manifold) weights " begin include("optimizers/structured_array_parameters.jl") end +@safetestset "map_to_cpu " begin + include("map_to_cpu_tests.jl") +end @info "Starting data and data-loader tests" @safetestset "Data " begin include("data/test_data.jl") diff --git a/test/transformer_related/multi_head_attention_stiefel_optim_cache.jl b/test/transformer_related/multi_head_attention_stiefel_optim_cache.jl index 7ba9c922a..932ca2ee9 100644 --- a/test/transformer_related/multi_head_attention_stiefel_optim_cache.jl +++ b/test/transformer_related/multi_head_attention_stiefel_optim_cache.jl @@ -23,7 +23,7 @@ function check_adam_cache(C::GeometricOptimizers.OptimizerCache{T}, tol=T(10) * _check_slahm_zero(C.δ, tol) _check_slahm_zero(C.m₁, tol) end -check_adam_cache(B::NamedTuple) = apply_toNT(check_adam_cache, B) +check_adam_cache(B::NamedTuple) = map(check_adam_cache, B) @doc raw""" This checks if the momentum cache was set up in the right way. @@ -34,7 +34,7 @@ function check_momentum_cache(C::GeometricOptimizers.OptimizerCache{T}, tol=T(10 @test C isa MomentumCache _check_slahm_zero(C.δ, tol) end -check_momentum_cache(B::NamedTuple) = apply_toNT(check_momentum_cache, B) +check_momentum_cache(B::NamedTuple) = map(check_momentum_cache, B) @doc raw""" This checks if the gradient cache was set up in the right way. @@ -44,7 +44,7 @@ function check_gradient_cache(C::GeometricOptimizers.OptimizerCache{T}) where T @test C isa GradientCache @test hasproperty(C, :δ) end -check_gradient_cache(B::NamedTuple) = apply_toNT(check_gradient_cache, B) +check_gradient_cache(B::NamedTuple) = map(check_gradient_cache, B) @doc raw""" This checks if all the caches are set up in the right way for the `MultiHeadAttention` layer with Stiefel weights. diff --git a/test/transformer_related/multi_head_attention_stiefel_retraction.jl b/test/transformer_related/multi_head_attention_stiefel_retraction.jl index b007dd66b..f113a3ce0 100644 --- a/test/transformer_related/multi_head_attention_stiefel_retraction.jl +++ b/test/transformer_related/multi_head_attention_stiefel_retraction.jl @@ -16,7 +16,7 @@ function check_retraction_geodesic(A::AbstractMatrix{T}, tol=eps(T)) where T @test typeof(A_retracted) <: StiefelManifold @test LinearAlgebra.norm(A_retracted - StiefelProjection(A_retracted)) < tol end -check_retraction_geodesic(cache::NamedTuple) = apply_toNT(check_retraction_geodesic, cache) +check_retraction_geodesic(cache::NamedTuple) = map(check_retraction_geodesic, cache) check_retraction_geodesic(B::MomentumCache) = check_retraction_geodesic(B.δ) @doc raw""" @@ -27,7 +27,7 @@ function check_retraction_cayley(A::AbstractMatrix{T}, tol=eps(T)) where T @test typeof(A_retracted) <: StiefelManifold @test LinearAlgebra.norm(A_retracted - StiefelProjection(A_retracted)) < tol end -check_retraction_cayley(cache::NamedTuple) = apply_toNT(check_retraction_cayley, cache) +check_retraction_cayley(cache::NamedTuple) = map(check_retraction_cayley, cache) check_retraction_cayley(B::MomentumCache) = check_retraction_cayley(B.δ) @doc raw""" diff --git a/test/transformer_related/multi_head_attention_stiefel_setup.jl b/test/transformer_related/multi_head_attention_stiefel_setup.jl index 181bfc866..52e8a19a6 100644 --- a/test/transformer_related/multi_head_attention_stiefel_setup.jl +++ b/test/transformer_related/multi_head_attention_stiefel_setup.jl @@ -12,7 +12,7 @@ function check_setup(A::AbstractMatrix{T}, tol=T(10)*eps(T)) where T @test typeof(A) <: StiefelManifold @test check(A) < tol end -check_setup(ps::NamedTuple) = apply_toNT(check_setup, ps) +check_setup(ps::NamedTuple) = map(check_setup, ps) check_setup(ps::NetworkParameters) = check_setup(GeometricMachineLearning.params(ps)) @doc raw""" @@ -22,7 +22,7 @@ function check_grad_setup(B::AbstractMatrix{T}, tol=T(10)*eps(T)) where T @test typeof(B) <: StiefelLieAlgHorMatrix @test LinearAlgebra.norm(B) < tol end -check_grad_setup(gx::NamedTuple) = apply_toNT(check_grad_setup, gx) +check_grad_setup(gx::NamedTuple) = map(check_grad_setup, gx) check_grad_setup(B::MomentumCache) = check_grad_setup(B.δ) @doc raw""" From 1cb3eb0dd141c893011e545a883df50c9088fe48 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Mon, 24 Aug 2026 11:56:51 +0900 Subject: [PATCH 2/3] Record 0.7.0 in the changelog --- CHANGELOG.md | 87 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 87 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0ca47c2d7..6cfdd0fe1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,93 @@ 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. +## [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` loses six per-type methods, `apply_toNT` turns out to have been +`Base.map` all along, and `_eltype` turns out to have been a less careful `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`, `_add` and `add!`'s `NamedTuple` method 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.** Not a rename — the + two disagree, and the disagreement is the point. `_eltype` returned the element type of the *first* + leaf and read a structured leaf's dense interface; `parameter_eltype` promotes across every leaf and + descends through `freeparameters`. So a layer mixing `Float32` and `Float64` weights used to pick + whichever came first in the `NamedTuple` and hand that to `Adam(T)` and to `_GMLGradient{T}`. + Unexported, so this is breaking only for code reaching into the package. + +### 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 six methods that + existed to unwrap and reconstruct a `StiefelManifold`, a `SymmetricMatrix`, a `SkewSymMatrix` and + the two triangular types collapse into one. `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. + `test/optimizers/structured_array_parameters.jl` is the regression net. + +### 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].** From b14326f68529793c2df56e82e4afc6f25e33eb2a Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Mon, 24 Aug 2026 16:45:18 +0900 Subject: [PATCH 3/3] Address review: the README, the section descent, and three counts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `README.md`'s headline example still called `GeometricMachineLearning.apply_toNT`, which this branch deletes -- the one surviving reference to the name outside the changelog. It is `map` there too. `_tree_optim_step!` had the same inversion this branch fixes in `_make_optimizer_cache`, one function further down: `λY isa NamedTuple ? λY[k] : λY` asks for one container type where the parameters are asked for two. Today only a `NamedTuple` arrives, because this package's own `GlobalSection(::NetworkParameters)` unwraps the container before `GeometricOptimizers` sees it -- and that method is the one it gives back next, its replacement upstream returning a container. Then a single `isa NamedTuple` hands every layer the whole section tree instead of its own section. The cache reordering had no test: `structured_array_parameters.jl` checks that training runs, which it does either way. `test/optimizers/utils/optimization_step.jl` now pins the *shape* -- a network's cache and state are `NamedTuple`s keyed by its layers, one `GeometricOptimizers` cache and one state per layer. Simulated against a widened `OptimizerSolution`, the old branch order does not merely give the root one cache: `OptimizerCache(::Adam, ::NetworkParameters)` is a `MethodError` outright. `_eltype` -> `parameter_eltype` was described as a fix, and it is not one. The two really do differ, but not at the four call sites this package had: `_eltype` was only ever asked for a `T` under `_use_go_cache`, and `OptimizerSolution{T}` is homogeneous in `T` by construction -- `ArrayNamedTuple{T} = NamedTuple{S,<:Tuple{Vararg{AbstractArray{T}}}}`. A layer mixing `Float32` and `Float64` weights 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 change buys is four fewer methods to own. Finally the two items the plan asked for and this branch had passed over. `add!`'s `NamedTuple` arm goes: nothing called it, and `AbstractNeuralNetworks.add!` is about a destination and two summands, which a parameter tree is not. `_add` and `add!` leave the export list -- `_add`'s siblings `_norm` and `_diff` were never on it and are the two of the three `src/` uses, and `add!` is reachable from the package that owns the generic. And `_add(::History, ::SingleHistory)`, which shared nothing but the name, is `_push_history!`, which also says that it mutates. --- CHANGELOG.md | 71 ++++++++++++++----- README.md | 2 +- src/GeometricMachineLearning.jl | 6 +- src/map_to_cpu.jl | 2 +- src/nnsolution/history.jl | 2 +- src/nnsolution/neural_net_solution.jl | 2 +- src/optimizers/optimizer.jl | 10 ++- src/utils.jl | 9 +-- test/optimizers/utils/optimization_step.jl | 24 ++++++- .../zygote_specific/dense_gpu.jl | 2 +- 10 files changed, 101 insertions(+), 29 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6cfdd0fe1..89f1115c5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,7 @@ 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. @@ -18,8 +18,8 @@ breaking release). **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` loses six per-type methods, `apply_toNT` turns out to have been -`Base.map` all along, and `_eltype` turns out to have been a less careful `parameter_eltype`. +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*. @@ -39,27 +39,54 @@ the half of this release with no visible effect today — see *Fixed*. beside an ordinary `Matrix` is no obstacle. Verified on Julia 1.10, the compat floor, as well as on 1.13. - `_norm`, `_diff`, `_add` and `add!`'s `NamedTuple` method 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 + `_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.** Not a rename — the - two disagree, and the disagreement is the point. `_eltype` returned the element type of the *first* - leaf and read a structured leaf's dense interface; `parameter_eltype` promotes across every leaf and - descends through `freeparameters`. So a layer mixing `Float32` and `Float64` weights used to pick - whichever came first in the `NamedTuple` and hand that to `Adam(T)` and to `_GMLGradient{T}`. +- **`_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 six methods that + 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. `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. + 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 @@ -88,7 +115,19 @@ the half of this release with no visible effect today — see *Fixed*. 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. - `test/optimizers/structured_array_parameters.jl` is the regression net. + + `_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 diff --git a/README.md b/README.md index dc693a09c..04891ed93 100644 --- a/README.md +++ b/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) diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index e319bcd7e..85544a7d4 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -172,8 +172,10 @@ include("activations/softmax.jl") # are these needed? export UnknownProblem, NothingFunction -# + operation has been overloaded to work with NamedTuples! -export _add, 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 diff --git a/src/map_to_cpu.jl b/src/map_to_cpu.jl index 3173c8226..dcfa97dff 100644 --- a/src/map_to_cpu.jl +++ b/src/map_to_cpu.jl @@ -2,7 +2,7 @@ # # 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 six per-type methods this replaces were doing by hand. +# 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. # diff --git a/src/nnsolution/history.jl b/src/nnsolution/history.jl index 943a538ec..403b2eab0 100644 --- a/src/nnsolution/history.jl +++ b/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 diff --git a/src/nnsolution/neural_net_solution.jl b/src/nnsolution/neural_net_solution.jl index 764821959..f8c69a732 100644 --- a/src/nnsolution/neural_net_solution.jl +++ b/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 diff --git a/src/optimizers/optimizer.jl b/src/optimizers/optimizer.jl index c41f5619e..796418593 100644 --- a/src/optimizers/optimizer.jl +++ b/src/optimizers/optimizer.jl @@ -281,12 +281,20 @@ end # 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 diff --git a/src/utils.jl b/src/utils.jl index bd8456e29..d07ec7b30 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -51,10 +51,11 @@ function add!(C::AbstractVecOrMat, A::AbstractVecOrMat, B::AbstractVecOrMat) C .= A + B end -# No caller in the package; kept because `add!` is `AbstractNeuralNetworks`' generic and this is the -# container arm of it, next to the `AbstractVecOrMat` base case above and the structured-type methods -# in `src/arrays/gml_extensions.jl`. -add!(dx₁::NamedTuple, dx₂::NamedTuple, dx₃::NamedTuple) = map(add!, dx₁, dx₂, dx₃) +# There used to be a `NamedTuple` arm of `add!` here, recursing with `apply_toNT`. Nothing in the +# package, the tests, the docs or the scripts ever called it, and `AbstractNeuralNetworks.add!` -- +# whose generic this is -- is about a destination and two summands, which a parameter *tree* is not. +# The `AbstractVecOrMat` base case above and the structured-type methods in +# `src/arrays/gml_extensions.jl` are the arms this package genuinely owns. # Type pyracy!! function Base.:+(a::Float64, b::Tuple{Float64}) diff --git a/test/optimizers/utils/optimization_step.jl b/test/optimizers/utils/optimization_step.jl index 3e3a8cd78..ec439f89c 100644 --- a/test/optimizers/utils/optimization_step.jl +++ b/test/optimizers/utils/optimization_step.jl @@ -1,5 +1,6 @@ using GeometricMachineLearning, Test, LinearAlgebra, KernelAbstractions using AbstractNeuralNetworks: AbstractExplicitLayer +import GeometricOptimizers import GeometricMachineLearning: NeuralNetwork import Random @@ -31,4 +32,25 @@ for N = 4:N_max for n = 1:N optimization_step_test(N, n, T) end -end \ No newline at end of file +end + +# The regression net for the branch order in `_make_optimizer_cache`/`_make_optimizer_state`: a +# `NetworkParameters` is a tree to descend into, so it is recognised *structurally*, before the +# capability question `_use_go_cache` asks. The shape that follows is one `GeometricOptimizers` cache +# per *layer* -- not one for the root, which is what putting the capability question first would give +# the moment `GeometricOptimizers` adds the container to `OptimizerSolution`, and not one per weight, +# which is what hoisting the `NamedTuple` branch as well would give. +@testset "one cache and one state per layer, keyed by the network's layers" begin + model = Chain(StiefelLayer(6, 3), Dense(6, 6, tanh)) + nn = NeuralNetwork(model, KernelAbstractions.CPU(), Float32) + o = Optimizer(AdamOptimizer(), nn) + + for tree in (o.cache, o.state) + @test tree isa NamedTuple + @test keys(tree) == keys(GeometricMachineLearning.params(nn)) + end + @test o.cache.L1 isa GeometricOptimizers.OptimizerCache + @test o.cache.L2 isa GeometricOptimizers.OptimizerCache + @test o.state.L1 isa GeometricOptimizers.OptimizerState + @test o.state.L2 isa GeometricOptimizers.OptimizerState +end diff --git a/test/performance_tests/zygote_specific/dense_gpu.jl b/test/performance_tests/zygote_specific/dense_gpu.jl index 833a322d5..4d8e7e781 100644 --- a/test/performance_tests/zygote_specific/dense_gpu.jl +++ b/test/performance_tests/zygote_specific/dense_gpu.jl @@ -34,7 +34,7 @@ function train(dev::Device, n_steps, batch_size) Base.Threads.@threads for i in 2:batch_size dat_in = data_input[batch[i]] dat_out = data_input[batch[i]] - dp = _add(dp, Zygote.gradient(ps -> loss(dat_in, dat_out, ps, st), ps)[1]) + dp = GeometricMachineLearning._add(dp, Zygote.gradient(ps -> loss(dat_in, dat_out, ps, st), ps)[1]) end GeometricMachineLearning.optimization_step!(opt, model, ps, dp)