diff --git a/CHANGELOG.md b/CHANGELOG.md index 0ca47c2d7..89f1115c5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,10 +9,136 @@ 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. +## [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` 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*. + +### 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` 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.** 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 five methods that + existed to unwrap and reconstruct a `StiefelManifold`, a `SymmetricMatrix`, a `SkewSymMatrix` and + 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 + 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. + + `_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 + +- `_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].** 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 5311be0c9..85544a7d4 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 @@ -166,8 +172,10 @@ include("activations/softmax.jl") # are these needed? export UnknownProblem, NothingFunction -# + operation has been overloaded to work with NamedTuples! -export _add, apply_toNT, 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 fc5a1a1c3..dcfa97dff 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 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. +# +# `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/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 5dbba5f17..796418593 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,13 +266,35 @@ 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. +# +# 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 d4d28a872..d07ec7b30 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,11 @@ 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 +# 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}) @@ -167,7 +162,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/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) 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"""