Take the remaining parameter walks from the packages that own them - #249
Conversation
`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.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #249 +/- ##
==========================================
+ Coverage 66.51% 67.02% +0.51%
==========================================
Files 99 99
Lines 2888 2869 -19
==========================================
+ Hits 1921 1923 +2
+ Misses 967 946 -21 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
`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.
michakraus
left a comment
michakraus
left a comment
There was a problem hiding this comment.
Reviewed the whole diff against the local environment (NNP 0.1.1, GO 0.4.2, ANN 0.7.0) and on Julia
1.10 and 1.13. The substance is right, and the two central claims hold up under direct checking
rather than assumption:
Base.mapoverNamedTuples does throwArgumentError("Named tuple names do not match.")on
mismatched, reordered and differently-sized keys, on 1.10 (the compat floor) as well as 1.13, and
the three-argument form maps as expected.apply_toNTreally wasBase.map.mapstoragereproduces the five structuredmap_to_cpumethods exactly. GO's
ext/NeuralNetworkParametersExt.jlsuppliesfreeparameters/rebuildfor all of them, and that
extension landed in GO 0.4.1 — which is this package's compat floor, so the floor holds without
a bump.NeuralNetworkParameters = "0.1"likewise coversmapstorageandparameter_eltype; both
are in 0.1.0.- The branch reordering is a no-op today, as claimed:
ArrayNamedTuple{T}is homogeneous inT, so a
NetworkParametersnever satisfied_use_go_cachein the first place.
The three failing nightly jobs are pre-existing — main fails the same way at
test/hamiltonian_neural_network_tests.jl:44, and the matrix marks nightly continue-on-error.
Six things came out of the review. All are fixed in b14326f, and the PR body is updated.
1. README.md:33 called GeometricMachineLearning.apply_toNT
The headline SympNet-on-GPU example, and the one surviving reference to the name outside the
changelog. Broken as of 0.7.0. It is map there now, with no package qualification needed.
2. _tree_optim_step! had the same inversion this PR fixes, one function further down
λY_k = λY isa NamedTuple ? λY[k] : λYOne container type where the parameters are asked for two. Today only a NamedTuple arrives, and the
reason it does is GeometricOptimizers.GlobalSection(ps::NetworkParameters) at the top of the same
file, which unwraps the container before GO sees it. That method is the one PR C hands back, and its
replacement upstream is expected to return a container of sections (plan A6: GlobalSectionNamedTuple
gets a sibling) — at which point isa NamedTuple is false and every layer is handed the whole section
tree instead of its own section. Same failure mode as the cache, same one-line prophylactic, and it
belongs in the PR whose stated purpose is forward-compatibility.
3. The headline fix had no test
structured_array_parameters.jl checks that training runs and the losses are finite, which is true
under either branch order. The cache shape is assertable today, so
test/optimizers/utils/optimization_step.jl now pins it: a network's cache and state are
NamedTuples keyed exactly by its layers, with one GeometricOptimizers cache and one state per
layer — not one for the root, and not one per weight.
I checked the counterfactual by simulating a widened OptimizerSolution (_use_go_cache extended to
accept NetworkParameters) against both branch orders. It is worse than the description says: under
the old order the root does not merely get one cache, it does not get one at all —
OptimizerCache(::Adam{Float32}, ::NetworkParameters{...}) is a MethodError outright. Under the new
order the same simulation still gives (:L1, :L2) with an AdamCache at each.
4. The _eltype → parameter_eltype story was wrong
The CHANGELOG, the PR body and the commit message all said a layer mixing Float32 and Float64
weights "used to pick whichever came first in the NamedTuple and hand that to Adam(T)". That
cannot happen. _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 reaching
_leaf_optim_step!/_go_update_leaf!, 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 mixed layer fails that
test and recurses to one cache per weight (each a GMLEuclideanState, since OptimizerSolution has no
AbstractMatrix arm), so _eltype never saw a mixed set.
The substitution is behaviour-preserving, and the better claim is the true one: parameter_eltype is
the upstream spelling, is strictly more careful, and agrees with _eltype at every call site this
package has because GO's homogeneity gate guarantees it. What the change buys is four fewer methods
to own. Rewritten in the CHANGELOG and the body; the paragraph in d4e85f3's message stays as it is,
since the branch is not being rewritten.
5. Two plan items were passed over without the PR saying so
Plan B4 asked for both, and both are done now:
add!(::NamedTuple, ::NamedTuple, ::NamedTuple)is deleted. Nothing called it, and
AbstractNeuralNetworks.add!— whose generic it was a method of — is about a destination and two
summands, which a parameter tree is not._addandadd!leave the export list, so that line goes entirely._add's siblings_normand
_diffwere never exported and are the two of the three thatsrc/reduced_system/actually calls;
add!is reachable from the package that owns the generic. Migration is one line
(using AbstractNeuralNetworks: add!) and is in the CHANGELOG._add(::History, ::SingleHistory), which shared nothing with the parameter-tree_addbut the name,
is_push_history!— which also records that it mutates its first argument, as the old name did not.
The third plan item, _GMLGradient(::NetworkParameters) and
_gml_rgrad(::NetworkParameters, ::NetworkParameters), is genuinely subsumed by the reordering: a
NetworkParameters can no longer reach a leaf under any version of GO, so both would be dead code.
Worth stating rather than leaving as a silent omission, which the body now does.
6. Counts and a stale note
"the six methods that existed to unwrap and reconstruct a StiefelManifold, a SymmetricMatrix, a
SkewSymMatrix and the two triangular types" names five; the sixth counted was the plain-array method,
which reconstructs nothing and survives as _to_host. Fixed in the CHANGELOG and in
src/map_to_cpu.jl's comment, and the arithmetic of the eight now closes. And the reconstruction note
at the top of the CHANGELOG still said "the 0.5.0 and 0.6.0 sections below, which were written
alongside the work" — 0.7.0 is one of them now.
No other findings. map is not shadowed anywhere in the package; the three
test/transformer_related/ conversions are all single-argument recursive helpers, where map and
apply_toNT agree exactly (zip of one NamedTuple yields 1-tuples, so fun(p...) was fun(value));
test/map_to_cpu_tests.jl passes as written, and parameter_eltype handles every leaf shape the
package can produce, Tuple included.
github-project-automation
Bot
moved this to Done
in Geometric Machine Learning Development
Finishes what 0.6.0 started. That release handed the HDF5 traversal of a parameter set over to
NeuralNetworkParameters and
GeometricOptimizers; this one does the same for the walks thatwere left, and makes the optimizer cache immune to a change
GeometricOptimizersis about to make.Nothing here needs an unreleased dependency — it builds against
NeuralNetworkParameters0.1.1 andGeometricOptimizers0.4.3, both in General — so it can merge on its own. It should merge beforethe
GeometricOptimizerscontainer change, for the reason in the last section.apply_toNTwasBase.mapmapoverNamedTuples takes any number of arguments and already throwsArgumentError: Named tuple names do not match.on mismatched or reordered keys — which is what thehand-rolled
@assert keys(ps[1]) == keys(p)was approximating, except that Base's check cannot becompiled out the way an
@assertcan. Heterogeneous values map fine, so aStiefelManifoldbeside anordinary
Matrixis no obstacle. Checked on Julia 1.10, the compat floor, as well as 1.13._norm,_diff,_addandadd!'s container arm callmapdirectly. That is the faithfultranslation rather than a simplification:
_diffand_addrecurse through their ownNamedTuplemethods, and
_normappliesnormone level down and divides by√length, so reaching for a deeperwalk there would change behaviour.
GeometricOptimizerscarries a character-identical copy of the same function, which this packagereached by qualified call. Dropping that call is what lets the copy go in its own release.
map_to_cpuis one walk instead of eight methodsNeuralNetworkParameters.mapstoragehands a function the storage of a leaf and rebuilds the leafaround the result, so the six methods that existed to unwrap and reconstruct a
StiefelManifold, aSymmetricMatrix, aSkewSymMatrixand the two triangular types collapse into one delegation.GeometricOptimizerssupplies the protocol for its own types, so nothing here knows which structuredtypes exist and one added upstream is covered without a change on this side.
mapstorageand notmapparameters: the latter hands the function whole leaves, which would stillneed a method per type to reach the storage. The
NeuralNetworkmethod stays — the docs tutorials andseveral scripts call it on a whole network.
It was untested, so the rewrite comes with
test/map_to_cpu_tests.jl._eltypebecomesparameter_eltypeThe two really are different functions —
_eltypereturned the element type of the first leaf andread a structured leaf's dense interface, where
parameter_eltypepromotes across every leaf anddescends through
freeparameters— but at the four call sites this package had they cannotdisagree, and an earlier draft of this description claimed a fix that cannot fire. Corrected.
_eltypewas only ever asked for aTunder_use_go_cache, which requiresx isa GeometricOptimizers.OptimizerSolution, and on theps_leafthat reaches_leaf_optim_step!,which is such an
x.OptimizerSolution{T}is homogeneous inTby construction — itsNamedTuplearm is
ArrayNamedTuple{T} = NamedTuple{S,<:Tuple{Vararg{AbstractArray{T}}}}. A layer mixingFloat32andFloat64weights therefore fails that test and recurses to one cache per weight, soa 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 where the walk is upstream's.
Why this should merge first
_make_optimizer_cacheand_make_optimizer_stateasked the capability question(
x isa GeometricOptimizers.OptimizerSolution) before the structural one, so aNetworkParametersreached the container branch only by virtue of not being in that union yet. The moment
GeometricOptimizersadopts the container, the root of a network would match instead: one cache forthe whole network rather than one per layer,
_leaf_optim_step!handed the entire tree, and_GMLGradienthanded aNetworkParametersit has no method for. AMethodErroron the first step ofevery training run, from a change that reads as purely additive upstream.
The
NetworkParametersbranch now comes first. TheNamedTuplebranch deliberately stays after thecapability test, and that asymmetry is the fix rather than an oversight: a layer is a
NamedTupleofarrays, which is exactly what one
GeometricOptimizerscache is for, so hoisting it too would descendinto the individual weights. Behaviour today is unchanged, which is what makes it safe to land
ahead of the upstream release rather than with it.
_tree_optim_step!had the same inversion one function further down, and it is fixed the same way:λY isa NamedTuple ? λY[k] : λYasked for one container type where the parameters are asked for two.Today only a
NamedTuplearrives, because this package's ownGlobalSection(::NetworkParameters)unwraps the container before
GeometricOptimizerssees it — and that method is the one GML gives backnext, its replacement upstream returning a container of sections. Then a single
isa NamedTuplehandsevery layer the whole section tree instead of its own section.
The reordering also had no test:
structured_array_parameters.jlchecks that training runs, whichit does either way.
test/optimizers/utils/optimization_step.jlnow pins the shape — a network'scache and state are
NamedTuples keyed by its layers, with oneGeometricOptimizerscache and onestate per layer. Simulated against a widened
OptimizerSolution, the old branch order turns out notmerely to give the root one cache:
OptimizerCache(::Adam, ::NetworkParameters)is aMethodErroroutright.
add!and_add, which the plan asked for and an earlier draft passed overadd!'sNamedTuplearm goes. Nothing in the package, the tests, the docs or the scripts called it,and
AbstractNeuralNetworks.add!— whose generic it was a method of — is about a destination and twosummands, which a parameter tree is not.
_addandadd!leave the export list, so that line goes entirely._add's siblings_normand_diffwere never exported and are the two of the three anything insrc/calls;add!is reachablefrom the package that owns the generic, this one only adding methods for the structured matrix types:
And
_add(::History, ::SingleHistory), which shared nothing with the parameter-tree_addbut thename, is
_push_history!— which also says that it mutates its first argument, as the old name didnot.
Not done, deliberately
_tree_optim_step!stays hand-written, and now says why. It walks the cache tree, which stops atthe layer where a cache sits, whereas
foreachparameterswalks the leaf protocol and would descendpast the layer into individual weights, re-pairing every cache with the wrong object. And
λYisbroadcast rather than zipped — a single
GlobalSectionmay stand in for a whole subtree — whichforeachparameterscannot express, since it takesvaluesof each trailing argument._GMLGradient(::NetworkParameters)and_gml_rgrad(::NetworkParameters, ::NetworkParameters), whichthe plan listed alongside the reordering, are subsumed by it: with the structural question asked
first, a
NetworkParametersalways recurses and can no longer arrive at a leaf, so neither method hasa caller in any version of
GeometricOptimizers. Adding them would be dead code.The
READMEexample is the one thing outsidesrc/that usedapply_toNT; it ismapthere now.Verification
Full suite green locally (56 blocks, zero failures), which is also what the
pre-pushhook ran.test/optimizers/utils/optimization_step.jlis the regression net for the cache reordering and pinsits shape;
test/optimizers/structured_array_parameters.jl— four architectures × four methods —covers the step itself.
test/map_to_cpu_tests.jlis new. The threetest/transformer_related/filesthat used
apply_toNTpass unchanged in substance.Base.map's key checking was verified directly rather than assumed: on 1.10 (the compat floor)and 1.13,
map(+, (a=1,b=2), (b=1,a=2))andmap(+, (a=1,b=2), (a=1,c=2))both throwArgumentError("Named tuple names do not match."), and the three-argument form maps as expected.The three failing nightly CI jobs are pre-existing:
mainfails the same way attest/hamiltonian_neural_network_tests.jl:44, and the matrix marks nightlycontinue-on-error.🤖 Generated with Claude Code