diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a84549a8..b4919d50b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,11 +17,75 @@ breaking release). ### Added +- **One optimizer method per leaf, through `CompositeMethod`.** `CompositeMethod` and + `CompositeState` are re-exported from `GeometricOptimizers`, with `ScalarMomentAdam` and + `ScalarMomentAdamState`. `Optimizer(CompositeMethod(; manifold = ScalarMomentAdam(), array = Adam()), nn)` + gives every leaf of the network's parameters the cache and the state of the method chosen for it, + which is what a network with Stiefel weights beside Euclidean ones needs to use `ScalarMomentAdam`: + that method steps a single `StiefelManifold`. A caller that wanted the split before had to add + methods to `_make_optimizer_cache`, `_make_optimizer_state` and `_leaf_optim_step!` from outside + this package. `test/optimizers/composite_method.jl` checks the per-leaf states and that a composite + of `Adam` on every leaf reproduces `Adam` exactly. - **The test suite runs on an Apple GPU in its `metal` group.** `test/devices/metal.jl` checks the two device paths that were measured by hand and not tested before: the products of a `PoissonTensor` with a wrapped device array, `𝕁 * view(A, [1, 2, 3, 4], :)` and `𝕁 * B'`, which `ext/GPUArraysCoreExt.jl` carries; and the `tensor_mat_mul!` kernel on `MetalBackend()`. Each result must stay on the device and match the host in `Float32`, with scalar indexing off. `Pkg.test(test_args = ["metal"])` runs them without the subject tests, and a missing device records a skip, which happens inside a sandbox. Metal (1.10 or later) is a test dependency on every platform; it installs and precompiles on Linux and Windows, and nothing there loads it. On an Apple M4 Max with Metal 1.11.1 all five assertions pass. - **A `Metal` workflow runs the Metal tests on GitHub's `macos-15` runner**, for the `min` and `1` Julia versions. The runner's GPU is Apple's paravirtual device, which Metal.jl supports from 1.10 and on macOS 15 or later. The workflow includes a step "Metal is functional" before the tests that fails the job where the device is unavailable, so it cannot pass without having run the tests. The job is not a required check. However, the Metal tests also run in two required checks—CI.yml's `Julia min - macOS-latest - default` and `Julia 1 - macOS-latest - default`—which execute `test/runtests.jl` with no test arguments on Apple silicon, selecting the `metal` group. - **`test/quality/jet.jl` checks the kernel launchers with JET.** It runs in `core`, directly after `quality/aqua.jl`, and takes about 14 s. Each function that launches a kernel of `src/` on the `CPU()` backend, the pullbacks among them, and the `Batch` functor that `test/data_loader/batch_index_set.jl` measures with `@allocated` get one `JET.report_opt` line per element type that the tests use, with `target_modules = (GeometricMachineLearning,)`. Each `@kernel` method also gets a line that analyses its generated `cpu_` function at the launcher's context, because JET drops the reports of a kernel body when it analyses the launcher. 142 lines pass on Julia 1.11.9 (JET 0.9.20) and 1.13.1 (JET 0.12.2). Three lines report a runtime dispatch on both versions and are `@test_broken` with issue #325: `PoissonTensor(::CPU, ::Int, ::DataType)`, and the input-and-output method of `convert_input_and_batch_indices_to_array` at `Float32` and `Float64`. Where JET does not work, the file records one `@test_skip`. JET joins `test/Project.toml` with no `[compat]` bound. `KNOWN_ISSUES.md` K2 records the Revise `EMFILE` messages that JET brings into the test log, and K3 the dispatches that the lines do not see: an unstable argument of a kernel launch, and two barriers in the Poisson-tensor kernel. - **A test calls `iterate` on `(q,)` initial conditions.** No test reached the `NeuralNetworkIntegrator` `iterate` method that takes `ics::NamedTuple{(:q,), Tuple{AT}}`. `test/docstrings/layers_and_architectures.jl` runs the docstring's `ResNet` from `(q = [1, 1, 1],)` for four points and expects the same trajectory as the vector form, `[1 2 4 8; 1 3 9 27; 1 3 7 15]`. +### Changed (breaking) + +- **The training step is `GeometricOptimizers.TrainingOptimizer`.** `Optimizer` holds one, over the + whole parameter set of the network, in its only field `training`; the fields `method`, `cache`, + `state`, `retraction`, `step_size` and `iterations` are gone. One cache and one state cover the + whole set, where there used to be one per layer, and every method steps every kind of leaf the same + way `GeometricOptimizers` does: no leaf is stepped by this package's own Euclidean update rules + (`GMLEuclideanState`) any more. `GeometricOptimizers`' G4 checked that `TrainingOptimizer` + reproduces this package's 0.8 iterates: exactly for `Adam`, and to `4eps(T)` for `GradientMethod` and + `MomentumMethod` on manifold weights. `ScalarMomentAdam` trains the Stiefel weights of a network + through a `CompositeMethod` (see *Added*); the step used to raise a `MethodError`, since this + package had no method of its own for it. +- **`optimization_step!(ps, opt, dp)` replaces `optimization_step!(opt, λY, ps, dp)`**, and + `optimize_for_one_epoch!` loses its last argument `λY`. The section is carried by the optimizer's + state, so the caller no longer builds a `GlobalSection` of the parameters for it. + `optimization_step!` is `GeometricOptimizers`' function, extended for `Optimizer`, so a `Main` that + loads both packages has one `optimization_step!` and not two. `dp` is the gradient a pullback + returns, a `NetworkParameters` or the bare `NamedTuple` of the same shape. +- **`retraction` takes `Cayley()` or `Geodesic()`**, the retraction types, and no longer the + functions `cayley` and `geodesic`. The default is `Cayley()`. +- **The methods carry no element type**, as in `GeometricOptimizers` 0.9: `Adam()`, + `Adam(; β₁, β₂, δ)`, `MomentumMethod(; α)` and `AdamOptimizerWithDecay(n_epochs; η₁, η₂)` in place + of `Adam(Float32)`, `Adam(T; …)`, `MomentumMethod(α)` and `AdamOptimizerWithDecay(n_epochs, T)`. + `Optimizer` converts the method to the element type of the parameters. +- **`LowerTriangular` and `UpperTriangular` are `StrictlyLowerTriangular` and + `StrictlyUpperTriangular`**, the names `GeometricOptimizers` 0.9 gives them, which no longer collide + with `LinearAlgebra`'s. A file saved with one of them by 0.8 or earlier is tagged with the old name + and still loads (JuliaGNI/GeometricOptimizers.jl#146); `test/integration/hdf5_support.jl` loads a + `VolumePreservingFeedForward` network written in the old layout. +- **Requires `GeometricOptimizers` 0.9, `AbstractNeuralNetworks` 0.9, `NeuralNetworkParameters` + 0.4.2, `GeometricBase` 0.15 and `SymbolicNeuralNetworks` 0.9.** No registered `GeometricOptimizers` + before 0.9 allows `NeuralNetworkParameters` 0.4 or `AbstractNeuralNetworks` 0.9, and no registered + `SymbolicNeuralNetworks` before 0.9 allows either (JuliaGNI/SymbolicNeuralNetworks.jl#80), so the + five move together. Under `NeuralNetworkParameters` 0.4 a Zygote gradient taken through a + `NetworkParameters` is the gradient with respect to each leaf's storage: a `SymmetricMatrix` leaf + keeps its type however often the loss reads it, which closes #319, and its off-diagonal entries are + twice those of the cotangent of the dense matrix. `test/helpers/network_parameters_gradient_structure.jl` + asserts both. + +### Removed (breaking) + +- **The three-axis methods of `Dense` and `Linear`, and the HDF5 extension.** `AbstractNeuralNetworks` + 0.9 defines both itself: `Dense`, `Linear` and `Affine` take inputs of rank 3 and more, and its + `HDF5Ext` has `save` and `load` for a `NeuralNetwork`, so `using HDF5` still loads them. Here they + were type piracy, the three methods in `src/layers/resnet.jl` and the four in `ext/HDF5Ext.jl`. + Aqua's `piracies` and `ambiguities` checks now run in `test/quality/aqua.jl` and report none; they + were switched off with counts of 3 and 1 in their place (*B7* in `KNOWN_ISSUES.md`, now closed). + `load(NeuralNetwork, …)` loses its `backend` keyword: the network loads on the CPU, and + `changebackend(backend, nn)` moves it. +- **`GMLEuclideanState`, `_GMLGradient` and the per-layer optimizer walk** (`_make_optimizer_cache`, + `_make_optimizer_state`, `_tree_optim_step!`, `_leaf_optim_step!`, `_as_go_solution`) are gone with + it, and with them `test/optimizers/gml_gradient_dispatch.jl` and + `test/optimizers/step_size_element_type.jl`, which tested two of them. *C2* in `KNOWN_ISSUES.md` + is closed; *C1* is narrowed to the one thing left, the name `Optimizer`, which both packages export. + ### Changed - **The package needs Julia 1.12 or later** (`[compat] julia = "1.12"`, was `"1.11"`). On Julia 1.11.9, the gradient of an `LNNLoss` with respect to the network parameters stops with `StackOverflowError` in type inference ("Internal error: during type inference of … Encountered stack overflow") on macOS and Windows; the `Julia min` job on Linux passes. The overflow comes from SymbolicUtils 4.49.0, whose `rrule` for `SymbolicUtils.Code.create_array` returns an exactly typed tuple of tangents, one per array element; Zygote's reverse pass through the generated Hessian of `lagrangian_acceleration` is then inferred far deeper. With the 4.48.0 rule restored, or in a task with a 64 MB stack, the same gradient passes on 1.11.9; Julia 1.12.7 and 1.13.1 pass it on the root task. This closes #327: the `Julia min` jobs on macOS and Windows had failed in the *Lagrangian Neural Network* testset since 2026-10-01, and every testset after it was never run there. diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index 7ebbc21f1..cccbe5b90 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -45,37 +45,6 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. every type parameter, and Aqua's `unbound_args` passes on Julia nightly. See *Fixed* above. The number is left vacant rather than reused.) -### B7 · Three methods are type piracy, and closing them is an API change rather than a deletion. - -- **location:** `src/layers/resnet.jl:63` -- **kind:** defect -- **found:** 2026-09-19 -- **evidence:** - Aqua's `piracies` check reports 3, down from the 12 this entry opened with; the other nine were - deleted in this release and are under *Removed (breaking)* above. All three are genuine under - Aqua's definition — the function and every argument type belong to other modules — and all three - are one thing: a functor on a *layer* type this package does not own, applied to a three-axis - array. - - | method | `src` | without GML | with GML | - |:--|:--|:--|:--| - | `(::Dense{M, N, true})(::AbstractArray{T, 3}, ::NamedTuple)` | `layers/resnet.jl:63` | `MethodError` | the layer applied along the third axis | - | `(::Dense{M, N, false})(::AbstractArray{T, 3}, ::NamedTuple)` | `layers/resnet.jl:67` | the same | the same | - | `(::Linear{M, N})(::AbstractArray{T, 3}, ::NamedTuple)` | `layers/resnet.jl:71` | the same | the same | - - `Dense` and `Linear` are `AbstractNeuralNetworks`', and upstream a 3-tensor argument is a - `MethodError` because `*` cannot take one. So the witness is not a changed answer but a new one: - a call that throws without this package loaded returns a value with it. - - Closing this means either `AbstractNeuralNetworks` gaining the three-axis methods, or this - package wrapping the two layer types in its own. Both are an API change, and neither is a - by-product of a piracy pass — which is why these three stayed when the other nine went. - - `src/layers/resnet.jl:63` is also the one remaining method ambiguity, against - `AbstractNeuralNetworks.Affine`. That pair is benign and has no witness at all: `Dense` is not a - subtype of `Affine` and the two have no common instance, so no call can reach it. `test/quality/aqua.jl` - asserts both counts — 3 piracies and 1 ambiguity — so neither can drift, in either direction. - ### B10 · `DataLoader(::EnsembleSolution{T, T1, Vector{ST}})` at `src/data_loader/data_loader.jl:325` has no test and appears unreachable through this package's current dependencies. - **location:** `src/data_loader/data_loader.jl:325` @@ -135,36 +104,18 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. ## C. Follow-ups and cleanups -### C1 · The parameter-tree traversal still belongs upstream. +### C1 · `Optimizer` is exported by both this package and GeometricOptimizers. -- **location:** — +- **location:** `src/optimizers/optimizer.jl` - **kind:** upstream - **found:** 2026-08-17 - **evidence:** - `_make_optimizer_cache`, - `_make_optimizer_state`, `_tree_optim_step!`, `_leaf_optim_step!` and the bespoke - `GMLEuclideanState` are GML implementations of what GeometricOptimizers supports natively for a - single parameter. `GMLEuclideanState` in particular duplicates what `GradientState`, - `MomentumState` and `AdamState` already do for a plain array. - - What has to go upstream is *not* a reuse of GeometricOptimizers' `Optimizer`: that one needs an - `OptimizerProblem`, i.e. an objective function, and minibatch training has none — the gradient - arrives from AD one batch at a time. It is a new entry point there, a - gradient-supplied-externally step over a `NamedTuple` parameter tree. GML's `Optimizer` would then - be the `NeuralNetwork` constructor and the training-loop functor, and nothing else. - - `Optimizer` is the one name still exported by both packages, so this is also what closes the last - of B1's class of collision. - -### C2 · Two `isa` branches remain in `_leaf_optim_step!` - -- **location:** — -- **kind:** defect -- **found:** 2026-08-16 -- **evidence:** - (for `AdamState`/`MomentumState`). - Measurement showed the traversal is not implicated in the compile-time problem, so this is tidying, - and it disappears entirely if C1 lands first. + The parameter-tree traversal this entry was opened for went upstream in 0.9: the training step is + `GeometricOptimizers.TrainingOptimizer`, and GML's `Optimizer` only pairs one with a network and + carries the training-loop functor. What remains is the name. GML's `Optimizer` and + GeometricOptimizers' `Optimizer` (the one `solve!` takes) are different types, both exported, so a + `Main` that `using`s both packages cannot write `Optimizer` unqualified. It is the last of B1's class + of collision. Closing it means renaming one of them, which is an API change in either package. ### C6 · Three generated MNIST PDFs are in this branch's history @@ -363,7 +314,7 @@ Not defects — claims this release makes that nothing has actually checked yet. ## Found late -### K1 · Four source comments still send the reader to the Open Issues section of `CHANGELOG.md`, which no longer exists. +### K1 · Two source comments still send the reader to the Open Issues section of `CHANGELOG.md`, which no longer exists. - **location:** `src/reduced_system/reduced_system.jl:109` - **kind:** found late @@ -373,17 +324,11 @@ Not defects — claims this release makes that nothing has actually checked yet. $ grep -rnE 'Open Issues|issue C1[48]|\*C21\*|\*B7\*' src test src/kernels/exponentials/tensor_exponential.jl:2:# name is wider than what it holds; renaming it, and `exponentials/` with it, is issue C18. src/reduced_system/reduced_system.jl:109:# issue C14 in `CHANGELOG.md`. - test/quality/aqua.jl:30:# layer types, and both are an API change rather than a tidy-up. They are *B7* under - test/quality/aqua.jl:31:# `## Open Issues` in `CHANGELOG.md` with that reasoning. test/reduced_system/reduced_system.jl:44:# for four of twelve, with four `NaN` and four in the opposite order. That is *C21* in - $ grep -n 'CHANGELOG' test/quality/aqua.jl - 31:# `## Open Issues` in `CHANGELOG.md` with that reasoning. - 63: # and it fails when one is removed without the entry above and in `CHANGELOG.md` going with it. ``` - The IDs are still valid; the comments at `src/reduced_system/reduced_system.jl:109`, - `test/quality/aqua.jl:31`, `test/quality/aqua.jl:63` and `test/reduced_system/reduced_system.jl:44–45` - name `CHANGELOG.md` rather than `KNOWN_ISSUES.md`. The fix is a follow-up change to those + The IDs are still valid; the comments at `src/reduced_system/reduced_system.jl:109` and + `test/reduced_system/reduced_system.jl:44–45` name `CHANGELOG.md` rather than `KNOWN_ISSUES.md`. The fix is a follow-up change to those comments. ## Upstream diff --git a/Project.toml b/Project.toml index 61360a273..2c969d945 100644 --- a/Project.toml +++ b/Project.toml @@ -25,30 +25,27 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [weakdeps] GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" -HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" [extensions] GPUArraysCoreExt = "GPUArraysCore" -HDF5Ext = "HDF5" [compat] -AbstractNeuralNetworks = "0.8" +AbstractNeuralNetworks = "0.9" ChainRulesCore = "1" ForwardDiff = "1" GPUArraysCore = "0.2" -GeometricBase = "0.14" +GeometricBase = "0.15" GeometricEquations = "0.21" -GeometricOptimizers = "0.8" +GeometricOptimizers = "0.9" GeometricSolutions = "0.6" -HDF5 = "0.16, 0.17" InteractiveUtils = "1" KernelAbstractions = "0.9" LinearAlgebra = "1" NNlib = "0.8, 0.9" -NeuralNetworkParameters = "0.3, 0.4" +NeuralNetworkParameters = "0.4.2" ProgressMeter = "1" Random = "1" -SymbolicNeuralNetworks = "0.8" +SymbolicNeuralNetworks = "0.9" Symbolics = "7" TimerOutputs = "0.5, 1" Zygote = "0.7" diff --git a/docs/check_references.jl b/docs/check_references.jl index e012f016f..a10e54245 100644 --- a/docs/check_references.jl +++ b/docs/check_references.jl @@ -12,14 +12,12 @@ # julia --project=docs docs/check_references.jl using GeometricMachineLearning -using HDF5 # so that the extension below is loaded, as in make.jl using AbstractNeuralNetworks # signatures in `@docs` blocks name types from here using Documenter using Documenter: DocSystem # keep in sync with `modules` in make.jl -const MODULES = Module[GeometricMachineLearning, - Base.get_extension(GeometricMachineLearning, :HDF5Ext)] +const MODULES = Module[GeometricMachineLearning] const SRC = joinpath(@__DIR__, "src") diff --git a/docs/make.jl b/docs/make.jl index 7309c5aa3..9fad1d468 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -4,6 +4,7 @@ include(joinpath(@__DIR__, "check_references.jl")) using GeometricMachineLearning +# loads `AbstractNeuralNetworks`' HDF5 extension, whose `save` and `load` the tutorials call using HDF5 using AbstractNeuralNetworks using Documenter @@ -387,8 +388,7 @@ makedocs(; # skips *sub*modules — and (ii) pull GeometricOptimizers' own docstrings in, whose internal # `@ref`s resolve in its namespace and point at bindings this manual does not document. # Names owned by GeometricOptimizers are referred to as plain code with a link to its manual. - modules = [ - GeometricMachineLearning, Base.get_extension(GeometricMachineLearning, :HDF5Ext)], + modules = [GeometricMachineLearning], authors = "Michael Kraus, Benedikt Brantner", repo = "https://github.com/JuliaGNI/GeometricMachineLearning.jl/blob/{commit}{path}#L{line}", sitename = "GeometricMachineLearning.jl", diff --git a/docs/src/architectures/abstract_neural_networks.md b/docs/src/architectures/abstract_neural_networks.md index 15a7cc19e..d6ea28316 100644 --- a/docs/src/architectures/abstract_neural_networks.md +++ b/docs/src/architectures/abstract_neural_networks.md @@ -47,11 +47,4 @@ and we see that it consists of two layers: a [`GradientLayerQ`](@ref) and a [`Gr ## Saving and Loading -`GeometricMachineLearning` adds HDF5-backed `save` and `load` methods for `NeuralNetwork` to the generics `NeuralNetworkParameters` defines. Writing and reading the parameter set itself belongs to that package; the structured parameter types (`StiefelManifold`, `SymmetricMatrix`, `SkewSymMatrix`, …) come back as themselves because `GeometricOptimizers`, which owns them, registers how each is rebuilt. Passing a prototype parameter set to `load` rebuilds against it and needs no registration at all. - -```@docs -save(::HDF5.H5DataStore, ::NeuralNetwork) -save(::AbstractString, ::NeuralNetwork) -load(::Type{NeuralNetwork}, ::HDF5.H5DataStore, ::AbstractNeuralNetworks.Architecture) -load(::Type{NeuralNetwork}, ::AbstractString, ::AbstractNeuralNetworks.Architecture) -``` +`save` and `load` for a `NeuralNetwork` are `AbstractNeuralNetworks`', which defines them in its HDF5 extension on the generics `NeuralNetworkParameters` provides: `using HDF5` loads them. Writing and reading the parameter set itself belongs to `NeuralNetworkParameters`; the structured parameter types (`StiefelManifold`, `SymmetricMatrix`, `SkewSymMatrix`, …) come back as themselves because `GeometricOptimizers`, which owns them, registers how each is rebuilt. Passing a prototype parameter set to `load` rebuilds against it and needs no registration at all. A network loads on the CPU; `changebackend(backend, nn)` moves it to a device. diff --git a/docs/src/arrays/tensors.md b/docs/src/arrays/tensors.md index a7e0d1eac..0a8d65a9d 100644 --- a/docs/src/arrays/tensors.md +++ b/docs/src/arrays/tensors.md @@ -1,6 +1,6 @@ # Tensors in `GeometricMachineLearning` -We typically store training data as *tensors with three axes* in `GeometricMachineLearning`. This allows for a parallel computation of matrix products, also for the special arrays such as [`LowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.LowerTriangular), [`UpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.UpperTriangular), [`SymmetricMatrix`](@extref GeometricOptimizers GeometricOptimizers.SymmetricMatrix) and [`SkewSymMatrix`](@extref GeometricOptimizers GeometricOptimizers.SkewSymMatrix) and objects of [`Manifold`](@extref GeometricOptimizers GeometricOptimizers.Manifold) type such as the [`StiefelManifold`](@extref GeometricOptimizers GeometricOptimizers.StiefelManifold). +We typically store training data as *tensors with three axes* in `GeometricMachineLearning`. This allows for a parallel computation of matrix products, also for the special arrays such as [`StrictlyLowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyLowerTriangular), [`StrictlyUpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyUpperTriangular), [`SymmetricMatrix`](@extref GeometricOptimizers GeometricOptimizers.SymmetricMatrix) and [`SkewSymMatrix`](@extref GeometricOptimizers GeometricOptimizers.SkewSymMatrix) and objects of [`Manifold`](@extref GeometricOptimizers GeometricOptimizers.Manifold) type such as the [`StiefelManifold`](@extref GeometricOptimizers GeometricOptimizers.StiefelManifold). ## Library Functions @@ -10,8 +10,8 @@ GeometricMachineLearning.tensor_mat_mul!(::AbstractArray{<:Number, 3}, ::Abstrac GeometricMachineLearning.tensor_mat_mul!(::AbstractArray{<:Number, 3}, ::AbstractArray{<:Number, 3}, ::SymmetricMatrix) GeometricMachineLearning.mat_tensor_mul(::AbstractMatrix, ::AbstractArray{<:Number, 3}) GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::AbstractMatrix, ::AbstractArray{<:Number, 3}) -GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::LowerTriangular, ::AbstractArray{<:Number, 3}) -GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::UpperTriangular, ::AbstractArray{<:Number, 3}) +GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::StrictlyLowerTriangular, ::AbstractArray{<:Number, 3}) +GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::StrictlyUpperTriangular, ::AbstractArray{<:Number, 3}) GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::SkewSymMatrix, ::AbstractArray{<:Number, 3}) GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::SymmetricMatrix, ::AbstractArray{<:Number, 3}) ``` \ No newline at end of file diff --git a/docs/src/index.md b/docs/src/index.md index 00d3dc8d6..7c9fe47f7 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -58,7 +58,7 @@ Everything tried on that device ran. Constructed, and then applied to both a mat Two of those results rest on `GeometricOptimizers` rather than on anything here. -- `VolumePreservingFeedForward` and `VolumePreservingTransformer` applied to a **matrix** multiply a `LowerTriangular` or `UpperTriangular` weight by that matrix. `GeometricOptimizers` supplies a `KernelAbstractions` kernel for that product, so the call does not fall through to a generic multiply that reads one entry at a time, which `GPUArraysCore` refuses. +- `VolumePreservingFeedForward` and `VolumePreservingTransformer` applied to a **matrix** multiply a `StrictlyLowerTriangular` or `StrictlyUpperTriangular` weight by that matrix. `GeometricOptimizers` supplies a `KernelAbstractions` kernel for that product, so the call does not fall through to a generic multiply that reads one entry at a time, which `GPUArraysCore` refuses. - `StiefelLayer`, `GrassmannLayer` and `PSDLayer`, and so every architecture that holds one of them, orthonormalize their weight at **construction** through `GeometricOptimizers.orthonormal_columns`. That is CholeskyQR2 — matrix products and triangular solves only — so it runs wherever the draw was allocated. `LinearAlgebra.qr!` is a host factorization and `Metal.jl` implements no `qr` for an `MtlArray`, so a device needs the other one. ## Tutorials diff --git a/docs/src/introduction.md b/docs/src/introduction.md index 7a9178980..08c516cd6 100644 --- a/docs/src/introduction.md +++ b/docs/src/introduction.md @@ -19,7 +19,7 @@ Closely linked to the research presented here is the development of a software p [^0]: This document was produced with `GeometricMachineLearning` `v0.3`. It may be that the interface will slightly change in future versions, but efforts will be made to keep these changes as small as possible. ```@docs; canonical = false -optimization_step! +optimization_step!(::NetworkParameters, ::Optimizer, ::Any) ``` So the docstring shows the name of the type or method, in most cases how to call it and then gives some information explaining what it does and potentially hyperlinks to other similar docstrings ([`Optimizer`](@ref) in this case); all of this information is indented by a tab. Docstrings may include other information under subheaders **Arguments** (showing the arguments the method can be supplied with), **Examples** (giving more detailed examples (including results) of how to use the method) and **Implementation** (giving details on how the method is implemented). When we reference a docstring it is always printed in blue (e.g. [`optimization_step!`](@ref)), indicating a hyperlink. In addition there is an *index of docstrings* showing all docstrings in chronological order with the associated page number. diff --git a/docs/src/layers/volume_preserving_feedforward.md b/docs/src/layers/volume_preserving_feedforward.md index d85cac9a8..e52bc26d6 100644 --- a/docs/src/layers/volume_preserving_feedforward.md +++ b/docs/src/layers/volume_preserving_feedforward.md @@ -7,7 +7,7 @@ The *volume-preserving feedforward layers* in `GeometricMachineLearning` are clo ```math \mathtt{VPFF}_{A, b}: x \mapsto x + \sigma(Ax + b), ``` -where ``\sigma`` is a nonlinearity, ``A`` is the weight and ``b`` is the bias. The matrix ``A`` is either a [`LowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.LowerTriangular) matrix ``L`` or an [`UpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.UpperTriangular) matrix ``U``. We demonstrate volume-preservation of these layers by considering the case ``A = L``. The matrix looks as follows: +where ``\sigma`` is a nonlinearity, ``A`` is the weight and ``b`` is the bias. The matrix ``A`` is either a [`StrictlyLowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyLowerTriangular) matrix ``L`` or a [`StrictlyUpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyUpperTriangular) matrix ``U``. We demonstrate volume-preservation of these layers by considering the case ``A = L``. The matrix looks as follows: ```math L = \begin{pmatrix} @@ -28,7 +28,7 @@ J = \nabla\mathtt{VPFF}_{L, b} = \begin{pmatrix} b_{n1} & \cdots & b_{n(n-1)} & 1 \end{pmatrix}, ``` -and the determinant of ``J`` is 1, i.e. the map is volume-preserving. A similar statement holds if the matrix ``A`` is [`UpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.UpperTriangular) instead of [`LowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.LowerTriangular). +and the determinant of ``J`` is 1, i.e. the map is volume-preserving. A similar statement holds if the matrix ``A`` is [`StrictlyUpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyUpperTriangular) instead of [`StrictlyLowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyLowerTriangular). ## Library Functions diff --git a/docs/src/optimizers/optimizer.md b/docs/src/optimizers/optimizer.md index ca0b095ee..bed0e14ce 100644 --- a/docs/src/optimizers/optimizer.md +++ b/docs/src/optimizers/optimizer.md @@ -6,18 +6,28 @@ global section — belongs to `GeometricOptimizers` and is described in its docu [Optimization on Homogeneous Spaces](@extref GeometricOptimizers :doc:`manifold_optimizers`) and [Retractions](@extref GeometricOptimizers :doc:`retractions`). -What `GeometricMachineLearning` adds is the part that is about *neural networks*: walking the -parameter tree of a `NeuralNetwork`, applying the right update to each leaf — a retraction -for a weight on a manifold, ordinary arithmetic for a Euclidean one — and driving that from a data -loader over epochs and batches. +The training step itself is `GeometricOptimizers`' `TrainingOptimizer` as well: one cache and one +state over the whole parameter set of a network, and a retraction for every weight on a manifold. +What `GeometricMachineLearning` adds is the part that is about *neural networks*: the +[`Optimizer`](@ref) that pairs a `TrainingOptimizer` with a `NeuralNetwork`, and the loop that drives +it from a data loader over epochs and batches. The gradient comes from automatic differentiation on one batch at a time, so there is no objective function to hand a line search; the step size is a property of the [`Optimizer`](@ref) instead. It is either a number, or a `GeometricOptimizers.DecayingStatic` schedule: ```julia -opt = Optimizer(Adam(Float32), nn; step_size = 1e-3) -opt = Optimizer(nn; AdamOptimizerWithDecay(n_epochs, Float32)...) +opt = Optimizer(Adam(), nn; step_size = 1e-3) +opt = Optimizer(nn; AdamOptimizerWithDecay(n_epochs)...) +``` + +A network whose weights are not all of one kind can take a different method for each kind. A +`CompositeMethod` chooses the method per leaf of the parameter set, and every leaf then gets the cache +and the state of the method chosen for it. `ScalarMomentAdam` steps a single `StiefelManifold`, so a +network with Stiefel weights beside Euclidean ones needs a composite to use it: + +```julia +opt = Optimizer(CompositeMethod(; manifold = ScalarMomentAdam(), array = Adam()), nn) ``` ## Library Functions @@ -25,6 +35,5 @@ opt = Optimizer(nn; AdamOptimizerWithDecay(n_epochs, Float32)...) ```@docs Optimizer optimize_for_one_epoch! -optimization_step! -GeometricMachineLearning._as_go_solution +optimization_step!(::NetworkParameters, ::Optimizer, ::Any) ``` diff --git a/docs/src/tutorials/grassmann_layer.md b/docs/src/tutorials/grassmann_layer.md index 1b1acbcd8..3c4eb27c6 100644 --- a/docs/src/tutorials/grassmann_layer.md +++ b/docs/src/tutorials/grassmann_layer.md @@ -105,13 +105,6 @@ nn = NeuralNetwork(model, CPU(), Float64) nothing # hide ``` -We then *lift* the neural network parameters via `GlobalSection`. - -```@example rosenbrock -λY = GlobalSection(params(nn)) -nothing # hide -``` - As the cost function ``c`` for the Wasserstein loss[^3] we simply use: [^3]: For each Wasserstein loss we need to define such a cost function. For this particular choice of ``c`` the Wasserstein distance corresponds to a ``W_2`` loss. @@ -284,7 +277,7 @@ loss_array = zeros(training_steps) for i in 1:training_steps val, dp = compute_gradient(params(nn)) loss_array[i] = val - optimization_step!(optimizer, λY, params(nn), dp) + optimization_step!(params(nn), optimizer, dp) end ``` diff --git a/docs/src/tutorials/hamiltonian_neural_network.md b/docs/src/tutorials/hamiltonian_neural_network.md index 4fe0b0914..5b2c8c42c 100644 --- a/docs/src/tutorials/hamiltonian_neural_network.md +++ b/docs/src/tutorials/hamiltonian_neural_network.md @@ -42,7 +42,7 @@ We can now train the network ```@example hnn batch = Batch(10) n_epochs = 100 -o = Optimizer(Adam(Float64), hnn) +o = Optimizer(Adam(), hnn) loss_array = o(hnn, dl, batch, n_epochs, loss) using CairoMakie # hide lines(loss_array) # hide @@ -84,7 +84,7 @@ nothing # hide We can now train the network: ```@example hnn -o_pairs = Optimizer(Adam(Float64), hnn_pairs) +o_pairs = Optimizer(Adam(), hnn_pairs) loss_array_pairs = o_pairs(hnn_pairs, dl_pairs, batch, n_epochs, loss_pairs) lines(loss_array_pairs) # hide ``` diff --git a/docs/src/tutorials/linear_symplectic_transformer.md b/docs/src/tutorials/linear_symplectic_transformer.md index 7d38af6d6..81010f9f8 100644 --- a/docs/src/tutorials/linear_symplectic_transformer.md +++ b/docs/src/tutorials/linear_symplectic_transformer.md @@ -63,7 +63,7 @@ nn_standard = NeuralNetwork(arch_standard) nn_symplectic = NeuralNetwork(arch_symplectic) nn_sympnet = NeuralNetwork(arch_sympnet) -o_pairing = AdamOptimizerWithDecay(n_epochs, Float64) +o_pairing = AdamOptimizerWithDecay(n_epochs) o_standard = Optimizer(nn_standard; o_pairing...) o_symplectic = Optimizer(nn_symplectic; o_pairing...) diff --git a/docs/src/tutorials/symplectic_autoencoder.md b/docs/src/tutorials/symplectic_autoencoder.md index ab958083c..3a206737e 100644 --- a/docs/src/tutorials/symplectic_autoencoder.md +++ b/docs/src/tutorials/symplectic_autoencoder.md @@ -145,7 +145,7 @@ dl = DataLoader(dl_cpu, backend, Float32) sae_nn_gpu = NeuralNetwork(sae_arch, CUDADevice(), Float32) -o = Optimizer(sae_nn_gpu; AdamOptimizerWithDecay(integrator_train_epochs, Float32)...) +o = Optimizer(sae_nn_gpu; AdamOptimizerWithDecay(integrator_train_epochs)...) # train the network o(sae_nn_gpu, dl, Batch(batch_size), n_epochs) diff --git a/docs/src/tutorials/symplectic_transformer.md b/docs/src/tutorials/symplectic_transformer.md index 75fcf178a..8d1331747 100644 --- a/docs/src/tutorials/symplectic_transformer.md +++ b/docs/src/tutorials/symplectic_transformer.md @@ -63,7 +63,7 @@ nn_standard = NeuralNetwork(arch_standard) nn_symplectic = NeuralNetwork(arch_symplectic) nn_sympnet = NeuralNetwork(arch_sympnet) -o_pairing = AdamOptimizerWithDecay(n_epochs, Float64) +o_pairing = AdamOptimizerWithDecay(n_epochs) o_standard = Optimizer(nn_standard; o_pairing...) o_symplectic = Optimizer(nn_symplectic; o_pairing...) diff --git a/docs/src/tutorials/sympnet_tutorial.md b/docs/src/tutorials/sympnet_tutorial.md index 602f58a2c..7d71eecae 100644 --- a/docs/src/tutorials/sympnet_tutorial.md +++ b/docs/src/tutorials/sympnet_tutorial.md @@ -100,7 +100,7 @@ We have to define an [optimizer](@extref GeometricOptimizers Standard-Neural-Net ```@example sympnet # set up optimizer; for this we first need to specify the optimization method -opt_method = Adam(type) +opt_method = Adam() # we then call the optimizer struct which allocates the cache la_opt = Optimizer(opt_method, la_nn) g_opt = Optimizer(opt_method, g_nn) diff --git a/docs/src/tutorials/volume_preserving_attention.md b/docs/src/tutorials/volume_preserving_attention.md index a4474faec..43abfb902 100644 --- a/docs/src/tutorials/volume_preserving_attention.md +++ b/docs/src/tutorials/volume_preserving_attention.md @@ -117,9 +117,9 @@ Next we train the networks (here we pick a batch size of 30 and train for 1000 e ```@setup volume_preserving_attention function set_up_optimizers(nn_skew, nn_arb, nn_comp) - o_skew = Optimizer(Adam(T), nn_skew) - o_arb = Optimizer(Adam(T), nn_arb) - o_comp = Optimizer(Adam(T), nn_comp) + o_skew = Optimizer(Adam(), nn_skew) + o_arb = Optimizer(Adam(), nn_arb) + o_comp = Optimizer(Adam(), nn_comp) o_skew, o_arb, o_comp end diff --git a/docs/src/tutorials/volume_preserving_transformer_rigid_body.md b/docs/src/tutorials/volume_preserving_transformer_rigid_body.md index 487c659c7..2d60ee879 100644 --- a/docs/src/tutorials/volume_preserving_transformer_rigid_body.md +++ b/docs/src/tutorials/volume_preserving_transformer_rigid_body.md @@ -199,7 +199,7 @@ const n_epochs = 500000 const batch_size = 16384 const feedforward_batch = Batch(batch_size) const transformer_batch = Batch(batch_size, seq_length, seq_length) -const opt_pairing = AdamOptimizerWithDecay(n_epochs, T; η₁ = 1e-2, η₂ = 1e-6) +const opt_pairing = AdamOptimizerWithDecay(n_epochs; η₁ = 1e-2, η₂ = 1e-6) o_vpff = Optimizer(nn_vpff; opt_pairing...) o_vpt = Optimizer(nn_vpt; opt_pairing...) diff --git a/ext/HDF5Ext.jl b/ext/HDF5Ext.jl deleted file mode 100644 index 39912a134..000000000 --- a/ext/HDF5Ext.jl +++ /dev/null @@ -1,87 +0,0 @@ -module HDF5Ext - -using HDF5 -using GeometricMachineLearning -import AbstractNeuralNetworks: NeuralNetworkBackend, Architecture -# `save`, `load`, `params` and the parameter container are `NeuralNetworkParameters`' as of -# `AbstractNeuralNetworks` 0.7, which only re-binds them; reach for them where they are defined. -import NeuralNetworkParameters: NetworkParameters, params, save, load - -# --------------------------------------------------------------------------- -# save / load — the entry points that dispatch on this package's `NeuralNetwork`. -# -# The traversal itself is not here. `NeuralNetworkParameters` walks the parameter set and records -# each group's key order; `GeometricOptimizers` says through `freeparameters`/`rebuild` where each -# structured matrix keeps its numbers, and registers the types so a file loads with no prototype. -# --------------------------------------------------------------------------- - -""" - save(h5::HDF5.H5DataStore, nn::NeuralNetwork) - -Save the parameters of `nn` into an already-open HDF5 store. - -Extends `save` with a dispatch on `NeuralNetwork`. The parameters themselves are written by -`NeuralNetworkParameters`, which tags each structured leaf with the type to rebuild it as and -records the key order of every group. -""" -save(h5::HDF5.H5DataStore, nn::NeuralNetwork) = save(h5, params(nn)) - -""" - save(filename::AbstractString, nn::NeuralNetwork) - -Convenience overload: open `filename` for writing, call [`save`](@ref) on the store, and return -`filename`. -""" -function save(filename::AbstractString, nn::NeuralNetwork) - HDF5.h5open(filename, "w") do h5 - save(h5, nn) - end - filename -end - -""" - load(::Type{NeuralNetwork}, h5::HDF5.H5DataStore, arch::Architecture; backend = CPU()) - load(::Type{NeuralNetwork}, h5::HDF5.H5DataStore, arch::Architecture, prototype; backend = CPU()) - -Load network parameters from an already-open HDF5 store and return a `NeuralNetwork` for `arch`. - -The element type is whatever the file holds, so a `Float32` network reloads as `Float32`. - -Structured parameters — `StiefelManifold`, `SymmetricMatrix` and the rest — are rebuilt from the -type each was stored under, which `GeometricOptimizers` registers with -`NeuralNetworkParameters.register_parameter_type!`. Pass `prototype`, a parameter set of the right -shape, to rebuild against it instead and skip the registry altogether. -""" -function load(::Type{NeuralNetwork}, h5::HDF5.H5DataStore, arch::Architecture; - backend::NeuralNetworkBackend = CPU()) - NeuralNetwork(arch, Chain(arch), load(NetworkParameters, h5), backend) -end - -function load(::Type{NeuralNetwork}, h5::HDF5.H5DataStore, arch::Architecture, prototype; - backend::NeuralNetworkBackend = CPU()) - NeuralNetwork(arch, Chain(arch), load(NetworkParameters, h5, prototype), backend) -end - -""" - load(::Type{NeuralNetwork}, filename::AbstractString, arch::Architecture; backend = CPU()) - load(::Type{NeuralNetwork}, filename::AbstractString, arch::Architecture, prototype; backend = CPU()) - -Convenience overload: open `filename` for reading, then call -[`load`](@ref) on the store. -""" -function load(::Type{NeuralNetwork}, filename::AbstractString, arch::Architecture; - backend::NeuralNetworkBackend = CPU()) - HDF5.h5open(filename, "r") do h5 - load(NeuralNetwork, h5, arch; backend = backend) - end -end - -function load( - ::Type{NeuralNetwork}, filename::AbstractString, arch::Architecture, prototype; - backend::NeuralNetworkBackend = CPU()) - HDF5.h5open(filename, "r") do h5 - load(NeuralNetwork, h5, arch, prototype; backend = backend) - end -end - -end diff --git a/scripts/reproduction/symplectic_autoencoders/training.jl b/scripts/reproduction/symplectic_autoencoders/training.jl index 389369520..064a0ca57 100644 --- a/scripts/reproduction/symplectic_autoencoders/training.jl +++ b/scripts/reproduction/symplectic_autoencoders/training.jl @@ -31,7 +31,7 @@ n_range = 2:1:smoke_size(15, 3) # The learning rate is the `Optimizer`'s `step_size`, not the method's: `Adam` lost the `η` field # it never applied to the direction, and `β₁`, `β₂` and `δ` became keywords so that the old # positional call fails instead of silently reading `η` as `β₁`. -opt = AdamOptimizer(T; β₁ = T(0.9), β₂ = T(0.99), δ = T(1.0e-8)) +opt = AdamOptimizer(; β₁ = T(0.9), β₂ = T(0.99), δ = T(1.0e-8)) const step_size = T(0.001) # The retraction belongs to the `Optimizer` now, not to the layer: `PSDLayer(M, N)` takes no # keyword. diff --git a/scripts/reproduction/sympnets/sympnet_pendulum.jl b/scripts/reproduction/sympnets/sympnet_pendulum.jl index 5fe8fa0f7..abc7cec98 100644 --- a/scripts/reproduction/sympnets/sympnet_pendulum.jl +++ b/scripts/reproduction/sympnets/sympnet_pendulum.jl @@ -70,14 +70,13 @@ method = AdamOptimizer() # initial gradients for calling Cache constructor opt = Optimizer(method, ps) # `optimization_step!` takes the global section of the parameters, not the model. -λY = GlobalSection(ps) # training println("initial loss: ", full_loss(ps, q, p)) training_steps = smoke_size(1000, 2) @showprogress for i in 1:training_steps dp = Zygote.gradient(ps -> loss(ps, q, p), ps)[1] - optimization_step!(opt, λY, ps, dp) + optimization_step!(ps, opt, dp) end println("final loss: ", full_loss(ps, q, p)) diff --git a/scripts/reproduction/volume_preserving_feedforward/double_pendulum.jl b/scripts/reproduction/volume_preserving_feedforward/double_pendulum.jl index 099bc56eb..3cb904625 100644 --- a/scripts/reproduction/volume_preserving_feedforward/double_pendulum.jl +++ b/scripts/reproduction/volume_preserving_feedforward/double_pendulum.jl @@ -42,7 +42,7 @@ const T = eltype(dl) # hyperparameters concerning training const n_epochs = smoke_size(1000, 2) const batch_size = 1024 -const opt_method = AdamOptimizer(T) +const opt_method = AdamOptimizer() const resnet_activation = tanh const t_validation = 30 diff --git a/scripts/reproduction/volume_preserving_transformer/rigid_body.jl b/scripts/reproduction/volume_preserving_transformer/rigid_body.jl index 47708f121..f84d93496 100644 --- a/scripts/reproduction/volume_preserving_transformer/rigid_body.jl +++ b/scripts/reproduction/volume_preserving_transformer/rigid_body.jl @@ -51,7 +51,7 @@ const dl = backend == CPU() ? DataLoader(T.(dl₁.input)) : # hyperparameters concerning training const n_epochs = smoke_size(500000, 2) const batch_size = 16384 -const opt_pairing = AdamOptimizerWithDecay(n_epochs, T; η₁ = 1e-2, η₂ = 1e-6) +const opt_pairing = AdamOptimizerWithDecay(n_epochs; η₁ = 1e-2, η₂ = 1e-6) # parameters for evaluation ics_val = [sin(1.1), 0.0, cos(1.1)] diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index bb474b69b..290c5d1d2 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -6,12 +6,11 @@ using AbstractNeuralNetworks # 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. +# One thing does 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`). +# A plain `NamedTuple` of layers is walked with `Base.map`, which needs nothing from anybody. import NeuralNetworkParameters: NetworkParameters -using NeuralNetworkParameters: mapparameters, mapstorage, parameter_eltype +using NeuralNetworkParameters: mapstorage using ChainRulesCore using GeometricBase using GeometricSolutions: GeometricSolution, EnsembleSolution, TimeSeries @@ -42,7 +41,7 @@ import Symbolics import GeometricOptimizers import GeometricOptimizers: Manifold, StiefelManifold, GrassmannManifold import GeometricOptimizers: SkewSymMatrix, SymmetricMatrix, - LowerTriangular, UpperTriangular, StiefelProjection + StrictlyLowerTriangular, StrictlyUpperTriangular, StiefelProjection import GeometricOptimizers: StiefelLieAlgHorMatrix, GrassmannLieAlgHorMatrix import GeometricOptimizers: rgrad, metric, check, global_section # `orthonormal_columns(draw)` orthonormalizes `draw()` with CholeskyQR2, which runs on whatever @@ -55,7 +54,13 @@ import GeometricOptimizers: Geodesic, Cayley, geodesic, cayley, retraction import GeometricOptimizers: OptimizerMethod, GradientMethod, MomentumMethod, Adam, GradientState, MomentumState, AdamState, - AdamOptimizerWithDecay, DecayingStatic + AdamOptimizerWithDecay, DecayingStatic, + ScalarMomentAdam, ScalarMomentAdamState, + CompositeMethod, CompositeState +# The training step is `GeometricOptimizers`' own: `Optimizer` holds a `TrainingOptimizer`, and +# `optimization_step!` gets a method for `Optimizer` rather than a second function of the name, so a +# `Main` that loads both packages has one `optimization_step!` and not two. +import GeometricOptimizers: TrainingOptimizer, optimization_step! import GeometricOptimizers: update! # `solve!` is imported rather than started afresh so that GML's `solve!(::NeuralNetwork{<:PSDArch}, # …)` — solve for the parameters directly, by SVD, instead of training for them — is a method of the @@ -114,7 +119,7 @@ include("arrays/poisson_tensor.jl") # Re-exported from GeometricOptimizers, so that `using GeometricMachineLearning` on its own still # gives a caller the matrix types its layers are parametrized by. export SymmetricMatrix, SkewSymMatrix -export LowerTriangular, UpperTriangular +export StrictlyLowerTriangular, StrictlyUpperTriangular export StiefelLieAlgHorMatrix, GrassmannLieAlgHorMatrix export StiefelProjection # GML's own @@ -201,6 +206,8 @@ export OptimizerMethod export GradientMethod, GradientState export MomentumMethod, MomentumState export Adam, AdamState +export ScalarMomentAdam, ScalarMomentAdamState +export CompositeMethod, CompositeState export Optimizer export optimization_step! export GlobalSection, global_section, apply_section, apply_section!, update_section! diff --git a/src/arrays/poisson_tensor.jl b/src/arrays/poisson_tensor.jl index 313276d65..4e7a1eb6d 100644 --- a/src/arrays/poisson_tensor.jl +++ b/src/arrays/poisson_tensor.jl @@ -39,6 +39,11 @@ struct PoissonTensor{T, AT <: AbstractMatrix{T}} <: AbstractMatrix{T} n::Int end +# The backend of the dense matrix it wraps. `GeometricOptimizers` asks it of both factors of a +# product with one of its manifolds or structured matrices, and refuses a product across two +# backends. +KernelAbstractions.get_backend(𝕁::PoissonTensor) = KernelAbstractions.get_backend(𝕁.J) + # The index types are named rather than left open. An untyped index pair claims every index type, # and it then collides with the `getindex` methods BandedMatrices and BlockArrays add to # `AbstractMatrix` for their own: nine ambiguities, of the same kind as the seventeen that the diff --git a/src/data_loader/optimize.jl b/src/data_loader/optimize.jl index 349448b9a..7cfe94fae 100644 --- a/src/data_loader/optimize.jl +++ b/src/data_loader/optimize.jl @@ -1,5 +1,5 @@ @doc raw""" - optimize_for_one_epoch!(opt, model, ps, dl, batch, loss, λY) + optimize_for_one_epoch!(opt, model, ps, dl, batch, loss) Sample the data contained in `dl` according to `batch` and optimize for these batches. @@ -24,7 +24,6 @@ All the arguments are mandatory (there are no defaults): 4. the data (i.e. an instance of [`DataLoader`](@ref)). 5. `batch`::[`Batch`](@ref): stores `batch_size` (and optionally `seq_length` and `prediction_window`). 6. `loss::NetworkLoss`. -7. the *section* `λY` of the parameters `ps`. # Implementation @@ -51,9 +50,8 @@ function optimize_for_one_epoch!(opt::Optimizer, model, ps::NetworkParameters, dl::DataLoader{T}, batch::Batch, - loss::NetworkLoss, - λY) where {T} - optimize_for_one_epoch!(opt, model, ps, dl, batch, ZygotePullback(loss), λY) + loss::NetworkLoss) where {T} + optimize_for_one_epoch!(opt, model, ps, dl, batch, ZygotePullback(loss)) end function optimize_for_one_epoch!(opt::Optimizer, @@ -61,8 +59,7 @@ function optimize_for_one_epoch!(opt::Optimizer, ps::NetworkParameters, dl::DataLoader{T}, batch::Batch, - _pullback::AbstractPullback, - λY) where {T} + _pullback::AbstractPullback) where {T} count = 0 total_error = T(0) batches = batch(dl) @@ -72,7 +69,7 @@ function optimize_for_one_epoch!(opt::Optimizer, loss_value, pullback = _pullback(ps, model, input_nt_output_nt) total_error += loss_value dp = _processing(pullback(one(loss_value))) - optimization_step!(opt, λY, ps, dp) + optimization_step!(ps, opt, dp) end total_error / count end @@ -83,13 +80,12 @@ function (o::Optimizer)(nn::NeuralNetwork, n_epochs::Integer, loss::NetworkLoss, _pullback::AbstractPullback = ZygotePullback(loss); show_progress = true) - Λ = GlobalSection(params(nn)) progress_object = show_progress ? ProgressMeter.Progress(n_epochs; enabled = true) : nothing loss_array = zeros(float(eltype(dl)), n_epochs) for i in 1:n_epochs loss_array[i] = optimize_for_one_epoch!( - o, nn.model, params(nn), dl, batch, _pullback, Λ) + o, nn.model, params(nn), dl, batch, _pullback) show_progress ? ProgressMeter.next!(progress_object; showvalues = [( :TrainingLoss, loss_array[i])]) : nothing diff --git a/src/kernels/kernel_ad_routines/mat_tensor_mul.jl b/src/kernels/kernel_ad_routines/mat_tensor_mul.jl index aa7970542..34742f6a8 100644 --- a/src/kernels/kernel_ad_routines/mat_tensor_mul.jl +++ b/src/kernels/kernel_ad_routines/mat_tensor_mul.jl @@ -73,15 +73,15 @@ function ChainRulesCore.rrule(::typeof(lo_mat_mul), S::AbstractVector{T}, return C, lo_mat_mul_pullback end -function ChainRulesCore.rrule(::typeof(mat_tensor_mul), B::LowerTriangular{T}, A::AbstractArray{ - T, 3}) where {T} +function ChainRulesCore.rrule(::typeof(mat_tensor_mul), B::StrictlyLowerTriangular{T}, + A::AbstractArray{T, 3}) where {T} @assert size(A, 1) == B.n C = mat_tensor_mul(B, A) function lower_triangular_mul_pullback(dC) dC = unthunk(dC) f̄, dS, dA, _ = rrule(lo_mat_mul, B.S, A, B.n)[2](dC) - return f̄, LowerTriangular(dS, B.n), dA + return f̄, StrictlyLowerTriangular(dS, B.n), dA end return C, lower_triangular_mul_pullback end @@ -144,15 +144,15 @@ function ChainRulesCore.rrule(::typeof(up_mat_mul), S::AbstractVector{T}, return C, up_mat_mul_pullback end -function ChainRulesCore.rrule(::typeof(mat_tensor_mul), B::UpperTriangular{T}, A::AbstractArray{ - T, 3}) where {T} +function ChainRulesCore.rrule(::typeof(mat_tensor_mul), B::StrictlyUpperTriangular{T}, + A::AbstractArray{T, 3}) where {T} @assert size(A, 1) == B.n C = mat_tensor_mul(B, A) function upper_triangular_mul_pullback(dC) dC = unthunk(dC) f̄, dS, dA, _ = rrule(up_mat_mul, B.S, A, B.n)[2](dC) - return f̄, UpperTriangular(dS, B.n), dA + return f̄, StrictlyUpperTriangular(dS, B.n), dA end return C, upper_triangular_mul_pullback end diff --git a/src/kernels/mat_tensor_mul.jl b/src/kernels/mat_tensor_mul.jl index a2bd3d1bc..e0490f465 100644 --- a/src/kernels/mat_tensor_mul.jl +++ b/src/kernels/mat_tensor_mul.jl @@ -119,7 +119,7 @@ function mat_tensor_mul!(C::AbstractArray{<:Number, 3}, A::SymmetricMatrix, B::A symmetric_mat_mul!(C, A.S, B, A.n) end -########################### LowerTriangular +########################### StrictlyLowerTriangular @kernel function lo_mul_kernel!(C::AbstractArray{T, 3}, S::AbstractVector{T}, B::AbstractArray{T, 3}, ::Int) where {T} @@ -152,21 +152,21 @@ function lo_mat_mul(S::AbstractVector{T}, B::AbstractArray{T, 3}, n::Int) where end @doc raw""" - mat_tensor_mul!(C, A::LowerTriangular, B) + mat_tensor_mul!(C, A::StrictlyLowerTriangular, B) Multiply the lower-triangular matrix `A` onto the tensor `B` from the left and store the result in `C`. This performs an efficient multiplication based on the special structure of the lower-triangular matrix `A`. """ -function mat_tensor_mul!(C::AbstractArray{<:Number, 3}, A::LowerTriangular, B::AbstractArray{ - <:Number, 3}) +function mat_tensor_mul!(C::AbstractArray{<:Number, 3}, A::StrictlyLowerTriangular, + B::AbstractArray{<:Number, 3}) @assert eltype(C) == eltype(A) == eltype(B) @assert A.n == size(C, 1) == size(B, 1) lo_mat_mul!(C, A.S, B, A.n) end -####################### UpperTriangular +####################### StrictlyUpperTriangular @kernel function up_mul_kernel!(C::AbstractArray{T, 3}, S::AbstractVector{T}, B::AbstractArray{T, 3}, n::Int) where {T} @@ -199,14 +199,14 @@ function up_mat_mul(S::AbstractVector{T}, B::AbstractArray{T, 3}, n::Int) where end @doc raw""" - mat_tensor_mul!(C, A::UpperTriangular, B) + mat_tensor_mul!(C, A::StrictlyUpperTriangular, B) Multiply the upper-triangular matrix `A` onto the tensor `B` from the left and store the result in `C`. This performs an efficient multiplication based on the special structure of the upper-triangular matrix `A`. """ -function mat_tensor_mul!(C::AbstractArray{<:Number, 3}, A::UpperTriangular, B::AbstractArray{ - <:Number, 3}) +function mat_tensor_mul!(C::AbstractArray{<:Number, 3}, A::StrictlyUpperTriangular, + B::AbstractArray{<:Number, 3}) @assert eltype(C) == eltype(A) == eltype(B) @assert A.n == size(C, 1) == size(B, 1) diff --git a/src/layers/resnet.jl b/src/layers/resnet.jl index c721c1e55..6a6b109b7 100644 --- a/src/layers/resnet.jl +++ b/src/layers/resnet.jl @@ -60,18 +60,6 @@ function (d::ResNetLayer{M, M, true})(x::AbstractArray{T, 3}, ps::NamedTuple) wh x + d.activation.(mat_tensor_mul(ps.weight, x) .+ ps.bias) end -function (d::Dense{M, N, true})(x::AbstractArray{T, 3}, ps::NamedTuple) where {M, N, T} - d.σ.(mat_tensor_mul(ps.W, x) .+ ps.b) -end - -function (d::Dense{M, N, false})(x::AbstractArray{T, 3}, ps::NamedTuple) where {M, N, T} - d.σ.(mat_tensor_mul(ps.W, x)) -end - -function (d::Linear{M, N})(x::AbstractArray{T, 3}, ps::NamedTuple) where {M, N, T} - mat_tensor_mul(ps.W, x) -end - function (d::Union{Dense{M, N}, ResNetLayer{M, N}})(z::QPT, ps::NamedTuple) where {M, N} @assert iseven(M) == iseven(N) == true @assert size(z.q, 1) * 2 == M diff --git a/src/layers/symplectic_attention.jl b/src/layers/symplectic_attention.jl index a3ba6d7c9..7d877d24d 100644 --- a/src/layers/symplectic_attention.jl +++ b/src/layers/symplectic_attention.jl @@ -128,11 +128,12 @@ function (d::SymplecticAttentionQ{M, M, :arbitrary})( p = z.p) end -# `A = ps.A` once, and then `A` twice: the local binding is what keeps the `SymmetricMatrix` -# structure in the Zygote gradient of this layer. Read `ps.A` twice and the gradient leaf degrades -# to a plain `Matrix` -- `test/parameters/symplectic_attention_network_parameters_gradient.jl` -# asserts both halves of that, and the last testset there covers this method. The binding costs no -# allocation; it is the number of accesses that decides the structure. +# `A = ps.A` once, and then `A` twice. Before `NeuralNetworkParameters` 0.4 the local binding was +# what kept the `SymmetricMatrix` structure in a Zygote gradient taken through a +# `NetworkParameters`: a second `ps.A` degraded the gradient leaf to a plain `Matrix` (issue #319). +# Since 0.4 that gradient is the storage gradient, which keeps the structure however often the leaf +# is read; `test/layers/symplectic_attention_network_parameters_gradient.jl` covers this method. The +# binding costs no allocation and stays. function (d::SymplecticAttentionQ{M, M, :symmetric})( z::NamedTuple{(:q, :p), Tuple{AT, AT}}, ps::NamedTuple) where {AT, M} A = ps.A diff --git a/src/layers/volume_preserving_feedforward.jl b/src/layers/volume_preserving_feedforward.jl index 7822fcda8..ef81bfc5e 100644 --- a/src/layers/volume_preserving_feedforward.jl +++ b/src/layers/volume_preserving_feedforward.jl @@ -4,10 +4,10 @@ Super-type of [`VolumePreservingLowerLayer`](@ref) and [`VolumePreservingUpperLayer`](@ref). The layers do the following: ```math -x \mapsto \begin{cases} \sigma(Lx + b) & \text{where $L$ is }\mathtt{LowerTriangular}, \\ \sigma(Ux + b) & \text{where $U$ is }\mathtt{UpperTriangular}. \end{cases} +x \mapsto \begin{cases} \sigma(Lx + b) & \text{where $L$ is }\mathtt{StrictlyLowerTriangular}, \\ \sigma(Ux + b) & \text{where $U$ is }\mathtt{StrictlyUpperTriangular}. \end{cases} ``` -The functor can be applied to a vector, a matrix or a tensor. The special matrices are implemented as [`LowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.LowerTriangular) and [`UpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.UpperTriangular). +The functor can be applied to a vector, a matrix or a tensor. The special matrices are implemented as [`StrictlyLowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyLowerTriangular) and [`StrictlyUpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyUpperTriangular). """ abstract type VolumePreservingFeedForwardLayer{M, N, bias} <: AbstractExplicitLayer{M, N} end @@ -36,7 +36,7 @@ x = \begin{pmatrix} 1 \\ 1 \end{pmatrix} using GeometricMachineLearning l = VolumePreservingLowerLayer(2, identity; use_bias=false) -A = LowerTriangular([1,], 2) +A = StrictlyLowerTriangular([1,], 2) ps = (weight = A, ) x = ones(eltype(A), 2) @@ -93,7 +93,7 @@ x = \begin{pmatrix} 1 \\ 1 \end{pmatrix} using GeometricMachineLearning l = VolumePreservingUpperLayer(2, identity; use_bias=false) -A = UpperTriangular([1,], 2) +A = StrictlyUpperTriangular([1,], 2) ps = (weight = A, ) x = ones(eltype(A), 2) @@ -141,7 +141,7 @@ function initialparameters( init_weight!(rng, S) init_bias!(rng, b) - (weight = LowerTriangular(S, M), bias = b) + (weight = StrictlyLowerTriangular(S, M), bias = b) end function initialparameters( @@ -153,7 +153,7 @@ function initialparameters( init_weight!(rng, S) init_bias!(rng, b) - (weight = UpperTriangular(S, M), bias = b) + (weight = StrictlyUpperTriangular(S, M), bias = b) end function initialparameters( @@ -163,7 +163,7 @@ function initialparameters( S = KernelAbstractions.allocate(backend, T, parameterlength(d)) init_weight!(rng, S) - (weight = LowerTriangular(S, M),) + (weight = StrictlyLowerTriangular(S, M),) end function initialparameters( @@ -173,7 +173,7 @@ function initialparameters( S = KernelAbstractions.allocate(backend, T, parameterlength(d)) init_weight!(rng, S) - (weight = UpperTriangular(S, M),) + (weight = StrictlyUpperTriangular(S, M),) end function (d::VolumePreservingFeedForwardLayer{M, M, :bias})( diff --git a/src/optimizers/optimizer.jl b/src/optimizers/optimizer.jl index 69d81f32a..886249676 100644 --- a/src/optimizers/optimizer.jl +++ b/src/optimizers/optimizer.jl @@ -1,423 +1,88 @@ -# Optimizer machinery on top of GeometricOptimizers. -# Kept out of `utils.jl` because it dispatches on `Manifold`, which is defined later. - -# Gradient wrapper: stores a pre-computed Euclidean gradient and applies rgrad on manifolds. -mutable struct _GMLGradient{T, VT} <: GeometricOptimizers.Gradient{T} - dp::VT -end - -_gml_rgrad(x::Manifold, dp) = rgrad(x, dp) -_gml_rgrad(x, dp) = dp -# `mapparameters` and not `map`: it recurses on the branches, so `_gml_rgrad` is only ever called on -# leaves, and it rebuilds in the shape of its *first* argument -- a container, which is what -# `GeometricOptimizers._copyto!(gradient_array(cache), ·)` has a method for. It also normalises its -# trailing arguments, so `dp` may stay the plain `NamedTuple` the pullback produced. -_gml_rgrad(x::NetworkParameters, dp) = mapparameters(_gml_rgrad, x, dp) - -(g::_GMLGradient{T})(x::NetworkParameters{T}) where {T} = _gml_rgrad(x, g.dp) -(g::_GMLGradient{T})(x::AbstractArray{T}) where {T} = g.dp - -# `Manifold` and not `AbstractArray` alone, because the two are otherwise **ambiguous** and not merely -# overlapping. `_GMLGradient{T}` is a `GeometricOptimizers.Gradient{T}`, `Manifold{T}` is an -# `AbstractMatrix{T}`, and `GeometricOptimizers` has `(::Gradient{T})(::Manifold{T})`; so on a bare -# manifold neither `(::_GMLGradient{T})(::AbstractArray{T})` above nor that one is more specific, and -# the call is a run-time `MethodError: ... is ambiguous`. This method is more specific than both. +# The optimizer of a training loop is `GeometricOptimizers.TrainingOptimizer`: one cache and one +# state over the whole parameter set, or one per leaf for a `CompositeMethod`, stepped with the +# gradient of one minibatch at a fixed or scheduled step size. What this package adds is the +# `Optimizer` that pairs one with a network, and the epoch loop in `data_loader/optimize.jl` that +# calls it. # -# The answer it gives is this package's: the gradient is already computed, so the projection is -# `rgrad(x, dp)` and not upstream's `rgrad(x, reshape(grad(vec(x)), ...))`, which would evaluate an -# inner gradient this functor does not have. `_gml_rgrad` is where that lives. -# -# `_make_optimizer_cache` below asks `_is_layer` first and so keeps a bare `Manifold` from arriving -# here at all. That ordering is a correctness fix in its own right -- one cache per layer rather than -# per weight -- but it must not be the *only* thing standing between a caller and an ambiguity. -(g::_GMLGradient{T})(x::Manifold{T}) where {T} = _gml_rgrad(x, g.dp) - -# A bare `AbstractVector`, for the same reason as the `Manifold` method above: `SimpleSolvers` has -# `(::Gradient{T})(::AbstractVector{T})`, `_GMLGradient{T} <: GeometricOptimizers.Gradient{T}`, and -# without a method of its own `(::_GMLGradient{T})(::AbstractArray{T})` above and that one are -# equally specific on a `Vector` -- the call is a run-time `MethodError: ... is ambiguous`. This -# method is more specific than both: `_GMLGradient{T} <: Gradient{T}` on the first argument and -# `AbstractVector{T} <: AbstractArray{T}` on the second. The answer is unchanged from the -# `AbstractArray{T}` method above -- the gradient is already Euclidean, so it is `g.dp`. -(g::_GMLGradient{T})(x::AbstractVector{T}) where {T} = g.dp +# There is no step here, no cache, no state and no list of methods. Each of those used to be written +# here a second time -- a per-layer walk, a Euclidean Adam of its own, a chain of `isa` tests over +# the state types -- and each could disagree with `GeometricOptimizers` without a test noticing. -# State for Euclidean (non-manifold) parameters. -mutable struct GMLEuclideanState{T, AT <: AbstractArray{T}} - iterations::Int - m₁::AT - m₂::AT -end -function GMLEuclideanState(x::AbstractArray{T}) where {T} - GMLEuclideanState{T, typeof(x)}(0, zero(x), zero(x)) -end - -# Adam with a decaying learning rate is GeometricOptimizers', and split the way it belongs: the -# direction is an `Adam` method, the schedule is a `DecayingStatic` line search. Both names are -# imported, and `Optimizer` below takes a `DecayingStatic` as its `step_size`. - -_is_go_native_method(::GeometricOptimizers.GradientMethod) = true -_is_go_native_method(::GeometricOptimizers.MomentumMethod) = true -_is_go_native_method(::GeometricOptimizers.Adam) = true -_is_go_native_method(::GeometricOptimizers.OptimizerMethod) = false - -function _adapt_method_to_T(method::GeometricOptimizers.Adam, ::Type{T}) where {T} - GeometricOptimizers.Adam(T; β₁ = T(method.β₁), β₂ = T(method.β₂), δ = T(method.δ)) -end -function _adapt_method_to_T(method::GeometricOptimizers.MomentumMethod, ::Type{T}) where {T} - GeometricOptimizers.MomentumMethod(T(method.α)) -end -_adapt_method_to_T(method, ::Type) = method - -# Whether one `GeometricOptimizers` cache covers `x` whole. -# -# A **layer** is a flat `NamedTuple` of arrays, and it is exactly what one cache is for. A `NamedTuple` -# whose values are branches is a *subtree* and gets one cache per layer instead, which is why the test -# is flatness and not merely "a `NamedTuple`". -# -# The rule lives here rather than upstream because a whole set of parameters reaches -# `GeometricOptimizers` only as a `NetworkParameters`, and a container says nothing about how deep the -# tree beneath it is. Splitting a network into layers is this package's decision, so this package makes -# it; [`_as_go_solution`](@ref) then does the wrap. -# -# A layer whose weights do not share an element type is still one cache, with `T` the promotion over -# them. `_adapt_method_to_T` reads that promotion, so the cache is built for the element type the -# layer actually has rather than for one its weights are required to agree on. -# `!isempty` and not `all` alone: `all` over an empty collection is vacuously true, so without it an -# empty set would be one cache's worth -- and `OptimizerCache` would then be asked for the element -# type of a set that has no leaves to promote. An empty set is zero caches' worth, so it descends and -# yields an empty tree of caches, which is what the step below iterates over zero times. -_is_layer(x::NamedTuple) = _all_leaves(values(x)) -_is_layer(x::NetworkParameters) = _all_leaves(values(x)) -_is_layer(_) = false - -_all_leaves(vs) = !isempty(vs) && all(v -> v isa AbstractArray, vs) - -function _use_go_cache(method, x) - _is_go_native_method(method) && - (x isa GeometricOptimizers.OptimizerSolution || _is_layer(x)) -end - -""" - _as_go_solution(x) - -`x` in the shape `GeometricOptimizers` takes a solution in. - -A layer given as a bare `NamedTuple` is wrapped; anything already in the right shape — a container, a -bare `Manifold`, an `AbstractArray` — is passed through. The wrap **shares the leaf arrays**, so an -in-place optimizer step writes through to the network's own weights and nothing has to be copied back. """ -_as_go_solution(x::NetworkParameters) = x -_as_go_solution(x) = _is_layer(x) ? NetworkParameters(x) : x + Optimizer(method, nn; retraction = Cayley(), step_size = default_step_size(method)) + Optimizer(nn; algorithm, linesearch = default_step_size(algorithm), retraction = Cayley()) -# 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`. -# -# That ordering is **load-bearing**. `NetworkParameters` *is* one of the types -# `GeometricOptimizers.OptimizerSolution` unions, so `_use_go_cache` is true at the root: without the -# test above, 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. -# -# One cache for the whole network is the better end state, and it is **not** what this does. -# `_GMLGradient` does take a `NetworkParameters`, but that is one wrapped *layer* and not the tree, -# because the container branch above splits the network first. Getting to one cache means dropping -# `_tree_optim_step!` and handing the whole tree over, which is a change of behaviour -- one -# `GlobalSection` tree and one `Q` across every layer -- so it wants its own release. Nothing here -# depends on it. -# -# **"Is this one cache's worth?" is asked before "is this a tree?"**, and a *flat* set answers yes to -# both. `_is_layer` is that question: a set whose values are all leaves is a layer, whether it arrives -# wrapped or as a branch, and one `GeometricOptimizers` cache is exactly what a layer is for. Asking -# the tree question first would descend into such a set and give every individual weight its own cache -# -- which for a manifold weight is not merely wasteful but wrong, since a bare `Manifold` then reaches -# the gradient functor where a whole layer should have. -# -# The tree branch comes second and covers both carriers, because a network is a tree of layers whether -# it is wrapped or not. -function _make_optimizer_cache(method, x) - if _is_go_native_method(method) && _is_layer(x) - GeometricOptimizers.OptimizerCache(_adapt_method_to_T(method, parameter_eltype(x)), - _as_go_solution(x)) - elseif x isa NetworkParameters || x isa NamedTuple - 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) - else - GMLEuclideanState(x) - end -end - -function _make_optimizer_state(method, x) - if _is_go_native_method(method) && _is_layer(x) - GeometricOptimizers.OptimizerState(_adapt_method_to_T(method, parameter_eltype(x)), - _as_go_solution(x)) - elseif x isa NetworkParameters || x isa NamedTuple - 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) - else - GMLEuclideanState(x) - end -end +The optimizer of a training loop for the parameters of `nn`, a `NeuralNetwork` or a +`NetworkParameters`. -""" - Optimizer(method, nn; retraction, step_size) - Optimizer(nn; algorithm, linesearch, retraction) +It holds one `GeometricOptimizers.TrainingOptimizer` over the whole parameter set, which +[`optimization_step!`](@ref) steps once per minibatch. `method` is a first-order method: +`GradientMethod()`, `MomentumMethod()`, `Adam()`, `ScalarMomentAdam()` or another member of +GeometricOptimizers' `AdamFamily`, or a `CompositeMethod`, which gives every leaf of the set the +method it selects for that leaf: -Optimizer state combining a `GeometricOptimizers` method with the parameters of a neural network. +```julia +method = CompositeMethod(; manifold = ScalarMomentAdam(), array = Adam()) +opt = Optimizer(method, nn) +``` -`step_size` is either a number — a fixed learning rate — or a -`GeometricOptimizers.DecayingStatic`, a learning rate that decays geometrically with the iteration -number. The second form is the one `GeometricOptimizers` uses itself, so a method paired with a -schedule splats straight in: +`step_size` is a number — a fixed learning rate — or a `DecayingStatic`, a learning rate that decays +geometrically with the iteration number. The keyword form takes the `(algorithm, linesearch)` pair +that `AdamOptimizerWithDecay` returns, so it splats straight in: ```julia -opt = Optimizer(nn; AdamOptimizerWithDecay(n_epochs, Float32)...) +opt = Optimizer(nn; AdamOptimizerWithDecay(n_epochs)...) ``` +`retraction` is `Cayley()` or `Geodesic()`. + # Extended help The step size is a property of the optimizer and not of the method: the same `Adam()` trains at any -learning rate. That is the split `GeometricOptimizers` makes — the method supplies a direction, a -`SimpleSolvers.LinesearchMethod` supplies how far to go along it — and `step_size` is GML's half of -it for a training loop, which has no objective function for a real line search to evaluate. +learning rate. That is the split `GeometricOptimizers` makes, where the method supplies a direction +and the step size says how far to go along it. A training loop has no objective for a line search to +evaluate, only the loss of one minibatch, so a step size is a number or a schedule and nothing else. """ -mutable struct Optimizer{MT <: GeometricOptimizers.OptimizerMethod, CT, ST, RT, SST} - method::MT - cache::CT - state::ST - retraction::RT - step_size::SST - iterations::Int +struct Optimizer{TO <: GeometricOptimizers.TrainingOptimizer} + training::TO end -_default_step_size(::GeometricOptimizers.Adam) = 1e-3 -_default_step_size(::GeometricOptimizers.OptimizerMethod) = 1e-2 - -_step_size(η::Real, ::Int) = Float64(η) -# `t` and not `t - 1`: `optimization_step!` increments before it asks, so the first step of a solve -# is `α(1) = γη₁`. That is what `DecayingStatic` means by iteration `t`: `solve!` calls -# `increase_iteration_number!` before `solver_step!`. -_step_size(ls::DecayingStatic, t::Int) = Float64(GeometricOptimizers.step_size(ls, t)) - -_current_step_size(opt::Optimizer, t::Int) = _step_size(opt.step_size, t) - -_optimizer_step_size(η::Real) = Float64(η) -# Anything that is not a plain number goes through the same funnel as the `linesearch` keyword below, -# so that the two entry points accept the same things: `step_size = Static(α)` is a fixed learning -# rate on both, and anything else reports the `ArgumentError` that explains why a real line search -# has nothing to search along here, instead of a `MethodError` naming this helper. -_optimizer_step_size(ls) = _step_size_from_linesearch(ls) - -function Optimizer(method::GeometricOptimizers.OptimizerMethod, nn::NeuralNetwork; - retraction = GeometricOptimizers.cayley, - step_size = _default_step_size(method)) - Optimizer(method, params(nn); retraction = retraction, step_size = step_size) +function Optimizer(method::OptimizerMethod, nn::NeuralNetwork; kwargs...) + Optimizer(method, params(nn); kwargs...) end -function Optimizer(method::GeometricOptimizers.OptimizerMethod, - ps::NetworkParameters; - retraction = GeometricOptimizers.cayley, - step_size = _default_step_size(method)) - Optimizer(method, _make_optimizer_cache(method, ps), _make_optimizer_state(method, ps), - retraction, _optimizer_step_size(step_size), 0) +function Optimizer(method::OptimizerMethod, ps::NetworkParameters; + retraction::GeometricOptimizers.AbstractRetraction = Cayley(), + step_size = GeometricOptimizers.default_step_size(method)) + Optimizer(TrainingOptimizer(ps; algorithm = method, linesearch = step_size, + retraction = retraction)) end -# The keyword form, so that the `(algorithm, linesearch)` pairing `GeometricOptimizers` returns from -# `AdamOptimizerWithDecay` splats in unchanged. `linesearch` is the step size under the name -# upstream gives it; a `Static` carries its own `α`, which is then the fixed learning rate. function Optimizer(nn_or_ps::Union{NeuralNetwork, NetworkParameters}; - algorithm::GeometricOptimizers.OptimizerMethod, - linesearch = nothing, - retraction = GeometricOptimizers.cayley, - step_size = linesearch === nothing ? _default_step_size(algorithm) : - _step_size_from_linesearch(linesearch)) - Optimizer(algorithm, nn_or_ps; retraction = retraction, step_size = step_size) -end - -_step_size_from_linesearch(ls::DecayingStatic) = ls -_step_size_from_linesearch(ls::GeometricOptimizers.Static) = Float64(ls.α) -function _step_size_from_linesearch(ls) - throw(ArgumentError( - "`Optimizer` takes a fixed step size or a `DecayingStatic` schedule, not a $(typeof(ls)). " * - "A training loop evaluates its loss on one batch at a time and has no objective for a line " * - "search to search along; use `GeometricOptimizers.Optimizer` with an `OptimizerProblem` for " * - "that.")) -end - -# Euclidean update rules -function _euclidean_update!(x::AbstractArray{T}, dx::AbstractArray, - state::GMLEuclideanState, ::GeometricOptimizers.GradientMethod, step_size) where {T} - x .-= T(step_size) .* dx -end -function _euclidean_update!(x::AbstractArray{T}, dx::AbstractArray, - state::GMLEuclideanState{T}, method::GeometricOptimizers.MomentumMethod, step_size) where {T} - # `p ← αp + ∇L`, the classic momentum recursion. The decay belongs on `p` and not on `∇L`: - # `p ← p + α∇L` is an undamped accumulator that grows without bound for a constant gradient - # instead of saturating at `∇L/(1 - α)`. Same recursion as GO's `update!(::MomentumState, ...)`. - state.m₁ .= T(method.α) .* state.m₁ .+ dx - x .-= T(step_size) .* state.m₁ -end -function _euclidean_update!(x::AbstractArray{T}, dx::AbstractArray, - state::GMLEuclideanState{T}, method::GeometricOptimizers.Adam, step_size) where {T} - t = state.iterations - _t = t + 1 - β₁, β₂, δ = T(method.β₁), T(method.β₂), T(method.δ) - # the first factor is `(β - β^t)/(1 - β^t)`, not `β/(1 - β^t)`: the latter is ~100x too large - # at t = 2 and compounds every step, which inflates the second moment until the update vanishes - fac₁₁ = (β₁-β₁^_t)/(1-β₁^_t) - fac₁₂ = (1-β₁)/(1-β₁^_t) - fac₂₁ = (β₂-β₂^_t)/(1-β₂^_t) - fac₂₂ = (1-β₂)/(1-β₂^_t) - state.m₁ .= fac₁₁ .* state.m₁ .+ fac₁₂ .* dx - state.m₂ .= fac₂₁ .* state.m₂ .+ fac₂₂ .* dx .^ 2 - x .-= T(step_size) .* state.m₁ ./ (sqrt.(state.m₂) .+ δ) -end -# `_euclidean_update!` has three methods and not four: Adam with a decaying learning rate *is* -# Adam — only `step_size` differs, and that comes in as an argument — so the `Adam` method above -# serves it too. - -function _go_update_leaf!(cache, state, local_grad, - method::GeometricOptimizers.Adam, ps_leaf) - GeometricOptimizers.update!(cache, state, local_grad, method, ps_leaf) -end - -function _go_update_leaf!(cache, state, local_grad, - method::GeometricOptimizers.MomentumMethod, ps_leaf) - GeometricOptimizers.update!(cache, state, local_grad, method, ps_leaf) -end - -function _go_update_leaf!(cache, state, local_grad, - method::GeometricOptimizers.OptimizerMethod, ps_leaf) - T = parameter_eltype(ps_leaf) - GeometricOptimizers.update!(cache, state, local_grad, - GeometricOptimizers.NoHessian{T}(), ps_leaf) -end - -# GO-managed leaf step (manifolds, vectors, whole layers) -# -# `ps` and not `ps_leaf` from here on: a layer is wrapped for the same reason its cache is, and the -# wrap shares the arrays, so `_copyto!(ps, solution(cache))` below writes into the weights the caller -# holds. A bare `Manifold` or `AbstractArray` passes through untouched. -function _leaf_optim_step!(cache::GeometricOptimizers.OptimizerCache, - state::GeometricOptimizers.OptimizerState, - dp_leaf, ps_leaf, λY_leaf, method, retraction, step_size) - T = parameter_eltype(ps_leaf) - ps = _as_go_solution(ps_leaf) - local_grad = _GMLGradient{T, typeof(dp_leaf)}(dp_leaf) - adapted = _adapt_method_to_T(method, T) - state.iterations += 1 - _go_update_leaf!(cache, state, local_grad, adapted, ps) - GeometricOptimizers._rmul!(GeometricOptimizers.direction(cache), T(step_size)) - GeometricOptimizers.update_section!(GeometricOptimizers.section(cache), - GeometricOptimizers.section(state), - GeometricOptimizers.direction(cache), - retraction) - GeometricOptimizers._copyto!(GeometricOptimizers.solution(cache), - GeometricOptimizers.section(cache)) - GeometricOptimizers._copyto!(ps, GeometricOptimizers.solution(cache)) - GeometricOptimizers._copyto!(λY_leaf, GeometricOptimizers.section(cache)) - # `section(cache)` is `update_section!(section(state), direction, retraction)`, so copying it is - # the same thing as retracting a second time -- and a retraction on a manifold is `O(N³)` where - # the copy is `O(N²)`. - GeometricOptimizers._copyto!(GeometricOptimizers.section(state), - GeometricOptimizers.section(cache)) - if state isa GeometricOptimizers.AdamState - GeometricOptimizers._copyto!(GeometricOptimizers.first_moment(state), - GeometricOptimizers.first_moment(cache)) - GeometricOptimizers._copyto!(GeometricOptimizers.second_moment(state), - GeometricOptimizers.second_moment(cache)) - elseif state isa GeometricOptimizers.MomentumState - # `p ← αp + ∇L`; see the note in `_euclidean_update!` for the momentum method. This has to - # match what `update!(::MomentumCache, ...)` anticipated when it formed the direction. - GeometricOptimizers._rmul!(GeometricOptimizers.momentum(state), adapted.α) - GeometricOptimizers._add!(GeometricOptimizers.momentum(state), - GeometricOptimizers.gradient_array(cache)) - end - nothing -end - -# Euclidean leaf step (plain AbstractArray params) -function _leaf_optim_step!(cache::GMLEuclideanState, state::GMLEuclideanState, - dp_leaf, ps_leaf, λY_leaf, method, retraction, step_size) - _euclidean_update!(ps_leaf, dp_leaf, state, method, step_size) - state.iterations += 1 - nothing -end - -# Whether `λY` is a tree of sections to descend into, or a single `GlobalSection` standing in for a -# whole subtree. Named rather than written as an `isa` union: it is one question, asked once, and the -# answer decides whether a layer gets its own section or the whole tree. A section tree is a plain -# `NamedTuple` today, because `GeometricOptimizers`' `GlobalSection(::NetworkParameters)` unwraps the -# container; the container arm is here so that a change there is a `MethodError` and not a silent -# hand-off of the whole tree to every layer. -_is_section_tree(::NamedTuple) = true -_is_section_tree(::NetworkParameters) = true -_is_section_tree(_) = false - -# 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. -# Only a `NamedTuple` ever arrives: `GeometricOptimizers`' own -# `GlobalSection(ps::NetworkParameters) = GlobalSection(params(ps))` unwraps the container, so a -# section tree is a plain `NamedTuple` whether the parameters are wrapped or not. Should that ever -# return a *container* of sections instead, `isa NamedTuple` alone would be false here 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 = _is_section_tree(λY) ? λY[k] : λY - _tree_optim_step!(caches[k], states[k], dp_k, ps[k], λY_k, - method, retraction, step_size) - end - else - _leaf_optim_step!(caches, states, dp, ps, λY, method, retraction, step_size) - end - nothing + algorithm::OptimizerMethod, + linesearch = GeometricOptimizers.default_step_size(algorithm), + retraction::GeometricOptimizers.AbstractRetraction = Cayley()) + Optimizer(algorithm, nn_or_ps; retraction = retraction, step_size = linesearch) end """ - optimization_step!(opt, λY, ps, dp) + optimization_step!(ps, opt::Optimizer, dp) -Apply one optimization step to the parameters `ps` and their gradient `dp`, with the method and step -size the [`Optimizer`](@ref) `opt` carries. +Take one step of the [`Optimizer`](@ref) `opt` with the gradient `dp` of the loss at the parameters +`ps`, and write the new parameters into `ps`. -`λY` is a `GlobalSection` of `ps` (or a `NamedTuple` of them). Note that it is an *output* -here: the section the optimizer carries from step to step lives in `opt.state`, and `λY` is written -so that callers who inspect it see the updated section. It therefore has to be allocated once and -reused, not rebuilt per step -- rebuilding it costs a QR decomposition per manifold weight. - -The step counter is incremented *before* the step size is read, so the first step of a run is step 1. -This matters for a decaying `step_size` and is how `GeometricOptimizers` counts too. +`ps` is the parameter set `opt` was built for. `dp` is the gradient a pullback returns, a +`NetworkParameters` or the bare `NamedTuple` of the same shape. This is +`GeometricOptimizers.optimization_step!` on the `TrainingOptimizer` that `opt` holds; see there for +what a step does and what it refuses. """ -function optimization_step!(opt::Optimizer, λY, ps, dp) - # The increment comes *first*, so the first step of a run is step 1. It matters only for a - # decaying `step_size`, and there it matters: reading the schedule before incrementing takes - # `α(0) = η₁`, one whole step above what `DecayingStatic` and `GeometricOptimizers.solve!` - # take. `solve!` counts the same way, by calling `increase_iteration_number!` before - # `solver_step!`. - opt.iterations += 1 - step = _current_step_size(opt, opt.iterations) - _tree_optim_step!(opt.cache, opt.state, dp, ps, λY, opt.method, opt.retraction, step) - nothing +function optimization_step!(ps::NetworkParameters, opt::Optimizer, dp) + optimization_step!(ps, opt.training, _as_parameter_set(dp)) + ps end +# A pullback seeded with a bare `NamedTuple` hands its gradient back as one; see `_get_contents`. +_as_parameter_set(dp::NetworkParameters) = dp +_as_parameter_set(dp::NamedTuple) = NetworkParameters(dp) + check(::Optimizer) = nothing diff --git a/test/arrays/triangular.jl b/test/arrays/triangular.jl index 06226a1b8..abb754f66 100644 --- a/test/arrays/triangular.jl +++ b/test/arrays/triangular.jl @@ -12,8 +12,8 @@ Random.seed!(1234) # third axis of a tensor with `mat_tensor_mul`, and the pullback of that kernel. function triangular_tensor_multiplication_test(T = Float64, n::Int = 5) - Aₗ = rand(LowerTriangular{T}, n) - Aᵤ = rand(UpperTriangular{T}, n) + Aₗ = rand(StrictlyLowerTriangular{T}, n) + Aᵤ = rand(StrictlyUpperTriangular{T}, n) B = rand(T, n, n, n) AₗB = mat_tensor_mul(Aₗ, B) @@ -25,8 +25,8 @@ function triangular_tensor_multiplication_test(T = Float64, n::Int = 5) end function triangular_tensor_multiplication_pullback_test(T = Float64, n::Int = 5) - Aₗ = rand(LowerTriangular{T}, n) - Aᵤ = rand(UpperTriangular{T}, n) + Aₗ = rand(StrictlyLowerTriangular{T}, n) + Aᵤ = rand(StrictlyUpperTriangular{T}, n) B = rand(T, n, n, n) C_diff = rand(T, n, n, n) diff --git a/test/data_loader/data_loader_optimization_step.jl b/test/data_loader/data_loader_optimization_step.jl index 9a9405bd0..65029a085 100644 --- a/test/data_loader/data_loader_optimization_step.jl +++ b/test/data_loader/data_loader_optimization_step.jl @@ -17,8 +17,7 @@ function test_data_loader(sys_dim, n_time_steps, n_params, T = Float32) dx = Zygote.gradient(ps -> loss(model, ps, dl.input, dl.input), ps)[1] ps_copy = deepcopy(ps) o = Optimizer(GradientOptimizer(), ps) - λY = GlobalSection(ps) - optimization_step!(o, λY, ps, dx) + optimization_step!(ps, o, dx) @test ps !== ps_copy end diff --git a/test/data_loader/optimizer_functor_with_adam.jl b/test/data_loader/optimizer_functor_with_adam.jl index ad449788e..b079c1609 100644 --- a/test/data_loader/optimizer_functor_with_adam.jl +++ b/test/data_loader/optimizer_functor_with_adam.jl @@ -50,8 +50,7 @@ function test_optimization_with_adam(; T = Float32, input_dim = 9, n_patches = 4 loss₁ = loss(model, ps, dl.input, dl.output) opt = Optimizer(Adam(), nn_obj) - λY = GlobalSection(ps) - loss_average = optimize_for_one_epoch!(opt, model, ps, dl, batch, loss, λY) + loss_average = optimize_for_one_epoch!(opt, model, ps, dl, batch, loss) loss₃ = loss(model, ps, dl.input, dl.output) diff --git a/test/helpers/network_parameters_gradient_structure.jl b/test/helpers/network_parameters_gradient_structure.jl index 114af9229..eda77a2a1 100644 --- a/test/helpers/network_parameters_gradient_structure.jl +++ b/test/helpers/network_parameters_gradient_structure.jl @@ -2,32 +2,25 @@ using Test using GeometricMachineLearning using GeometricMachineLearning: params using LinearAlgebra: norm +using NeuralNetworkParameters: storage_gradient using Zygote: gradient """ -Assert which shape the `SymmetricMatrix` leaf of a gradient has, and that the shape is cosmetic. +Assert the shape and the value of the `SymmetricMatrix` leaf of a gradient taken with respect to a +`NetworkParameters` wrapper. -A gradient taken with respect to a `NetworkParameters` wrapper keeps its `SymmetricMatrix` leaf -only when the loss expression performs a single `getproperty` on that wrapper. A second access -degrades the leaf to a plain `Matrix`. What decides the split is the number of `getproperty` calls -on the wrapper itself -- not how often the leaf is used, and not how often the plain `NamedTuple`s -below the wrapper are accessed. `p -> (A = p.L1.A; sum(A) + sum(A))` uses the leaf twice through one -access and keeps the structure; `p -> sum(p.L1.A) + sum(p.L1.A)` loses it. A `Zygote.gradient` taken -with respect to the bare wrapped `NamedTuple` (`params(ps)`) keeps the structure at any access -count, because `gradient` projects its result and `ChainRulesCore.ProjectTo` restores the leaf for a -`NamedTuple` but is `identity` for a `NetworkParameters`. The underlying `Zygote.pullback` loses the -structure on the second access either way. `CHANGELOG.md` records the experiment. +Since `NeuralNetworkParameters` 0.4 that gradient is the gradient with respect to the leaf's +*storage*: `NeuralNetworkParameters` converts the cotangent of each leaf with `storage_gradient`. A +`SymmetricMatrix` keeps each off-diagonal number once and shows it twice, so its storage gradient +doubles the off-diagonal entries of the cotangent of the dense matrix, and it is a `SymmetricMatrix` +however often the loss reads the leaf. Until 0.4 a second `getproperty` on the wrapper degraded the +leaf to a plain `Matrix` (issue #319). -A wrapper unwrapped with `values` rather than `getproperty` -- which is how `Chain` hands a layer -its parameters -- counts differently: there the accesses on the plain `NamedTuple` below decide. -See the layer testset in -`test/layers/symplectic_attention_network_parameters_gradient.jl`. - -`preserves_structure` states which side of that split `f` sits on. Both gradients are asserted to -agree numerically either way, so losing the structure is a change of type and not of value. +A `Zygote.gradient` taken with respect to the bare wrapped `NamedTuple` (`params(ps)`) is the +cotangent of the dense matrix, projected back onto a `SymmetricMatrix` by +`ChainRulesCore.ProjectTo`. The two are asserted to differ by exactly `storage_gradient`. """ -function network_parameters_gradient_structure_test( - f, ps::NetworkParameters, preserves_structure::Bool) +function network_parameters_gradient_structure_test(f, ps::NetworkParameters) g_wrapped = gradient(_ps -> norm(f(_ps)), ps)[1] g_bare = gradient(_ps -> norm(f(_ps)), params(ps))[1] @@ -35,11 +28,6 @@ function network_parameters_gradient_structure_test( a_bare = g_bare.L1.A @test typeof(a_bare) <: SymmetricMatrix - if preserves_structure - @test typeof(a_wrapped) <: SymmetricMatrix - else - @test_broken typeof(a_wrapped) <: SymmetricMatrix # issue #319 - end - - @test isapprox(Matrix(a_wrapped), Matrix(a_bare)) + @test typeof(a_wrapped) <: SymmetricMatrix + @test isapprox(Matrix(a_wrapped), Matrix(storage_gradient(ps.L1.A, a_bare))) end diff --git a/test/integration/changebackend_tests.jl b/test/integration/changebackend_tests.jl index e816364bf..19c23bf3b 100644 --- a/test/integration/changebackend_tests.jl +++ b/test/integration/changebackend_tests.jl @@ -35,8 +35,8 @@ leaves = ( grassmann = rand(GrassmannManifold{Float64}, N, n), symmetric = SymmetricMatrix(rand(10), 4), skew = SkewSymMatrix(rand(6), 4), - lower = LowerTriangular(rand(6), 4), - upper = UpperTriangular(rand(6), 4), + lower = StrictlyLowerTriangular(rand(6), 4), + upper = StrictlyUpperTriangular(rand(6), 4), stiefhor = StiefelLieAlgHorMatrix(SkewSymMatrix(rand(n, n)), rand(N - n, n), N, n), grasshor = GrassmannLieAlgHorMatrix(rand(N - n, n), N, n) ) diff --git a/test/integration/hdf5_support.jl b/test/integration/hdf5_support.jl index 6e04ea5f8..7e29aa89f 100644 --- a/test/integration/hdf5_support.jl +++ b/test/integration/hdf5_support.jl @@ -14,8 +14,8 @@ _ps_eq(a::AbstractArray, b::AbstractArray) = a ≈ b _ps_eq(a::StiefelManifold, b::StiefelManifold) = a.A ≈ b.A _ps_eq(a::SymmetricMatrix, b::SymmetricMatrix) = a.S ≈ b.S && a.n == b.n _ps_eq(a::SkewSymMatrix, b::SkewSymMatrix) = a.S ≈ b.S && a.n == b.n -_ps_eq(a::LowerTriangular, b::LowerTriangular) = a.S ≈ b.S && a.n == b.n -_ps_eq(a::UpperTriangular, b::UpperTriangular) = a.S ≈ b.S && a.n == b.n +_ps_eq(a::StrictlyLowerTriangular, b::StrictlyLowerTriangular) = a.S ≈ b.S && a.n == b.n +_ps_eq(a::StrictlyUpperTriangular, b::StrictlyUpperTriangular) = a.S ≈ b.S && a.n == b.n function _ps_eq(a::NamedTuple, b::NamedTuple) Set(keys(a)) == Set(keys(b)) || return false all(_ps_eq(a[k], b[k]) for k in keys(a)) @@ -48,7 +48,8 @@ function _write_legacy(h5, Y::StiefelManifold, path::AbstractString) end for (T, name) in ((:SymmetricMatrix, "SymmetricMatrix"), (:SkewSymMatrix, "SkewSymMatrix"), - (:LowerTriangular, "LowerTriangular"), (:UpperTriangular, "UpperTriangular")) + # the tags GML 0.8 and earlier wrote, under the names the two types had then + (:StrictlyLowerTriangular, "LowerTriangular"), (:StrictlyUpperTriangular, "UpperTriangular")) @eval function _write_legacy(h5, A::$T, path::AbstractString) g = HDF5.create_group(h5, path) HDF5.attributes(g)["gml_type"] = $name @@ -131,8 +132,8 @@ end end end -# VolumePreservingFeedForward uses LowerTriangular / UpperTriangular parameters. -@testset "save/load roundtrip: VolumePreservingFeedForward (LowerTriangular/UpperTriangular)" begin +# VolumePreservingFeedForward uses StrictlyLowerTriangular / StrictlyUpperTriangular parameters. +@testset "save/load roundtrip: VolumePreservingFeedForward (the triangular matrices)" begin arch = VolumePreservingFeedForward(4, 4, 1, tanh) nn = NeuralNetwork(arch) x = rand(4) @@ -215,7 +216,10 @@ end @testset "a file in the old gml_type layout still loads" begin for (name, arch, dimin) in (("LASympNet (SymmetricMatrix)", LASympNet(4), 4), ("SymplecticAutoencoder (StiefelManifold)", - SymplecticAutoencoder(10, 4), 10)) + SymplecticAutoencoder(10, 4), 10), + # tagged `LowerTriangular`/`UpperTriangular`, the names before GeometricOptimizers 0.9 + ("VolumePreservingFeedForward (StrictlyLowerTriangular, StrictlyUpperTriangular)", + VolumePreservingFeedForward(4, 4, 1, tanh), 4)) @testset "$name" begin nn = NeuralNetwork(arch) ps = params(nn) diff --git a/test/layers/double_multiplication_network_parameters_gradient.jl b/test/layers/double_multiplication_network_parameters_gradient.jl index d2fb3557c..0079a62dc 100644 --- a/test/layers/double_multiplication_network_parameters_gradient.jl +++ b/test/layers/double_multiplication_network_parameters_gradient.jl @@ -22,12 +22,12 @@ ps = NetworkParameters((L1 = (A = S,),)) t = rand(4, 4) @testset "_custom_mul: NetworkParameters gradient structure" begin - @testset "single_multiplication (one wrapper access: structure kept)" begin + @testset "single_multiplication" begin network_parameters_gradient_structure_test( - _ps -> single_multiplication(t, _ps), ps, true) + _ps -> single_multiplication(t, _ps), ps) end - @testset "double_multiplication (two wrapper accesses: structure lost)" begin + @testset "double_multiplication" begin network_parameters_gradient_structure_test( - _ps -> double_multiplication(t, _ps), ps, false) + _ps -> double_multiplication(t, _ps), ps) end end diff --git a/test/layers/gradient_layer_tests.jl b/test/layers/gradient_layer_tests.jl index f0eb7474a..622e6dc1c 100644 --- a/test/layers/gradient_layer_tests.jl +++ b/test/layers/gradient_layer_tests.jl @@ -24,19 +24,18 @@ end function test_gradient_layer_derivative_and_update(T, M, N, batch_size = 10) dummy_model = Chain(GradientLayerP(M, N, tanh), GradientLayerQ(M, N, tanh)) ps = params(NeuralNetwork(dummy_model, CPU(), T)) - o = Optimizer(Adam(T; β₁ = T(0.9), β₂ = T(0.999), δ = T(3e-7)), ps; + o = Optimizer(Adam(; β₁ = T(0.9), β₂ = T(0.999), δ = T(3e-7)), ps; step_size = T(0.1)) # test for vector x = rand(T, M) gs = Zygote.gradient(ps -> sum(dummy_model(x, ps)), ps)[1] - λY = GlobalSection(ps) - optimization_step!(o, λY, ps, gs) + optimization_step!(ps, o, gs) # test for matrix X = rand(T, M, batch_size) gs = Zygote.gradient(ps -> sum(dummy_model(X, ps)), ps)[1] - optimization_step!(o, λY, ps, gs) + optimization_step!(ps, o, gs) end types = (Float32, Float64) diff --git a/test/layers/manifold_layers.jl b/test/layers/manifold_layers.jl index aa51054e5..bc5d7b032 100644 --- a/test/layers/manifold_layers.jl +++ b/test/layers/manifold_layers.jl @@ -6,13 +6,12 @@ Random.seed!(1234) function stiefel_layer_test(T, M, N, tol = 1.0f-1) model = Chain(StiefelLayer(M, N), StiefelLayer(N, N)) ps = NeuralNetwork(model, T).params - o = Optimizer(Adam(T; β₁ = T(5.0f-1), β₂ = T(5.0f-1), δ = T(3.0f-7)), ps; + o = Optimizer(Adam(; β₁ = T(5.0f-1), β₂ = T(5.0f-1), δ = T(3.0f-7)), ps; step_size = T(1.0f0)) dx = (L1 = (weight = rand(T, N, M),), L2 = (weight = rand(T, N, N),)) ps_copy = deepcopy(ps) - λY = GlobalSection(ps) - optimization_step!(o, λY, ps, dx) + optimization_step!(ps, o, dx) # check that the new weight is different from the old one @test norm(ps_copy[1].weight - ps[1].weight) > T(tol) # check that the updated elements are on the Stiefel Manifold @@ -23,14 +22,13 @@ end function grassmann_layer_test(T, M, N, tol = 1.0f-1) model = Chain(GrassmannLayer(M, N), StiefelLayer(N, N)) ps = NeuralNetwork(model, T).params - o = Optimizer(Adam(T; β₁ = T(5.0f-1), β₂ = T(5.0f-1), δ = T(3.0f-7)), ps; + o = Optimizer(Adam(; β₁ = T(5.0f-1), β₂ = T(5.0f-1), δ = T(3.0f-7)), ps; step_size = T(1.0f0)) dx = (L1 = (weight = rand(T, N, M),), L2 = (weight = rand(T, N, N),)) ps_copy = deepcopy(ps) - λY = GlobalSection(ps) for i in 1:4 - optimization_step!(o, λY, ps, dx) + optimization_step!(ps, o, dx) end # check that the new weight is different from the old one @test norm(ps_copy[1].weight - ps[1].weight) > T(tol) diff --git a/test/layers/multi_head_attention_stiefel_optim_cache.jl b/test/layers/multi_head_attention_stiefel_optim_cache.jl index c94822e9a..8a0a29ef5 100644 --- a/test/layers/multi_head_attention_stiefel_optim_cache.jl +++ b/test/layers/multi_head_attention_stiefel_optim_cache.jl @@ -6,9 +6,9 @@ using GeometricMachineLearning, Test import GeometricOptimizers using GeometricOptimizers: GradientCache, MomentumCache, AdamCache import Random, LinearAlgebra -# The cache and the state hold a `NetworkParameters` per layer: a layer is wrapped at the -# `GeometricOptimizers` boundary, and the wrap shares the leaf arrays. `mapparameters` and not `map`, -# because it recurses on the branches and so reaches the leaves whichever shape it is handed. +# The cache holds one `NetworkParameters` over the whole network, the shape of the parameters. +# `mapparameters` and not `map`, because it recurses on the branches and so reaches the leaves +# whichever shape it is handed. # # Each helper below has two methods, because two shapes arrive: a wrapped layer is a # `NetworkParameters` and the tree above it is a plain `NamedTuple`. @@ -73,9 +73,9 @@ function test_cache_setups_for_optimizer_for_multihead_attention_layer(T::Type, o₂ = Optimizer(MomentumMethod(), ps) o₃ = Optimizer(GradientMethod(), ps) - check_adam_cache(o₁.cache) - check_momentum_cache(o₂.cache) - check_gradient_cache(o₃.cache) + check_adam_cache(o₁.training.cache) + check_momentum_cache(o₂.training.cache) + check_gradient_cache(o₃.training.cache) end test_cache_setups_for_optimizer_for_multihead_attention_layer(Float32, 64, 8) diff --git a/test/layers/multi_head_attention_stiefel_retraction.jl b/test/layers/multi_head_attention_stiefel_retraction.jl index 7b38a33ca..9b340e6b9 100644 --- a/test/layers/multi_head_attention_stiefel_retraction.jl +++ b/test/layers/multi_head_attention_stiefel_retraction.jl @@ -53,7 +53,7 @@ function test_multi_head_attention_retraction( model = Chain(MultiHeadAttention(dim, n_heads, Stiefel = true)) ps = NeuralNetwork(model, backend, T).params - cache = Optimizer(MomentumMethod(), ps).cache + cache = Optimizer(MomentumMethod(), ps).training.cache check_retraction_geodesic(cache) diff --git a/test/layers/multi_head_attention_stiefel_setup.jl b/test/layers/multi_head_attention_stiefel_setup.jl index 5a90ea42c..20a55d59c 100644 --- a/test/layers/multi_head_attention_stiefel_setup.jl +++ b/test/layers/multi_head_attention_stiefel_setup.jl @@ -50,7 +50,7 @@ function check_multi_head_attention_stiefel_setup(T::Type, N::Int, n::Int) @test nleaves > 0 @test check_setup_calls[] == nleaves - gx = Optimizer(MomentumMethod(), ps).cache + gx = Optimizer(MomentumMethod(), ps).training.cache check_grad_setup(gx) end diff --git a/test/layers/symplectic_attention_network_parameters_gradient.jl b/test/layers/symplectic_attention_network_parameters_gradient.jl index 405794832..518fa56bc 100644 --- a/test/layers/symplectic_attention_network_parameters_gradient.jl +++ b/test/layers/symplectic_attention_network_parameters_gradient.jl @@ -28,28 +28,24 @@ ps = NetworkParameters((L1 = (A = S,),)) t = (q = rand(2, 2), p = rand(2, 2)) @testset "Symplectic attention: NetworkParameters gradient structure" begin - @testset "symplectic_attention (two wrapper accesses: structure lost)" begin + @testset "symplectic_attention" begin network_parameters_gradient_structure_test( - _ps -> symplectic_attention(t, _ps), ps, false) + _ps -> symplectic_attention(t, _ps), ps) end - @testset "symplectic_attention_simplified (one wrapper access: structure kept)" begin + @testset "symplectic_attention_simplified" begin network_parameters_gradient_structure_test( - _ps -> symplectic_attention_simplified(t, _ps), ps, true) + _ps -> symplectic_attention_simplified(t, _ps), ps) end - @testset "symplectic_linear_map (one wrapper access: structure kept)" begin + @testset "symplectic_linear_map" begin network_parameters_gradient_structure_test( - _ps -> symplectic_linear_map(t, _ps), ps, true) + _ps -> symplectic_linear_map(t, _ps), ps) end end -# The layer keeps the structure although its body reads `A` twice. `Chain` reaches the layer -# through `values(_ps)` rather than through a `getproperty` on the wrapper, and on that route the -# count that decides the split is the one on the plain `NamedTuple` the layer is handed. The local -# binding at `src/layers/symplectic_attention.jl` makes that count one, and is what keeps the -# structure here: the same body written without it, reached the same way, degrades the leaf to a -# plain `Matrix`. +# The same through a layer, which `Chain` reaches through `values(_ps)` rather than through a +# `getproperty` on the wrapper. @testset "SymplecticAttentionQ layer: NetworkParameters gradient structure" begin l = SymplecticAttentionQ(4; symmetric = true) nn = NeuralNetwork(Chain(l)) - network_parameters_gradient_structure_test(_ps -> nn(t, _ps), params(nn), true) + network_parameters_gradient_structure_test(_ps -> nn(t, _ps), params(nn)) end diff --git a/test/layers/transformer_optimizer.jl b/test/layers/transformer_optimizer.jl index 1d7ffcfd7..9597b15dc 100644 --- a/test/layers/transformer_optimizer.jl +++ b/test/layers/transformer_optimizer.jl @@ -26,12 +26,9 @@ function transformer_gradient_test(T, dim, n_heads, L, seq_length = 8, batch_siz ps₂ = deepcopy(ps) ps₃ = deepcopy(ps) - λY₁ = GlobalSection(ps₁) - λY₂ = GlobalSection(ps₂) - λY₃ = GlobalSection(ps₃) - optimization_step!(o₁, λY₁, ps₁, dx) - optimization_step!(o₂, λY₂, ps₂, dx) - optimization_step!(o₃, λY₃, ps₃, dx) + optimization_step!(ps₁, o₁, dx) + optimization_step!(ps₂, o₂, dx) + optimization_step!(ps₃, o₃, dx) @test typeof(ps₁) == typeof(ps₂) == typeof(ps₃) == typeof(ps) @test ps₁[1].PQ.head_1 ≈ ps₂[1].PQ.head_1 @test ps₁[1].PQ.head_1 ≉ ps₃[1].PQ.head_1 diff --git a/test/layers/volume_preserving_feedforward.jl b/test/layers/volume_preserving_feedforward.jl index 673951afc..b8f50cbb4 100644 --- a/test/layers/volume_preserving_feedforward.jl +++ b/test/layers/volume_preserving_feedforward.jl @@ -35,11 +35,11 @@ test_volume_preserving_feedforward(; T = Float32) test_volume_preserving_feedforward(; T = Float64) # A vector in, a vector out. Neither layer shapes the result: the `:no_bias` method is -# `x + d.activation.(ps.weight * x)`, and `LowerTriangular * ::AbstractVector` returns a vector, -# which is what `LinearAlgebra` returns for every other matrix type. +# `x + d.activation.(ps.weight * x)`, and `StrictlyLowerTriangular * ::AbstractVector` returns a +# vector, which is what `LinearAlgebra` returns for every other matrix type. l = VolumePreservingLowerLayer(2, identity; use_bias = false) -A = LowerTriangular([1], 2) +A = StrictlyLowerTriangular([1], 2) @test l(ones(eltype(A), 2), (weight = A,)) == [1, 2] l = VolumePreservingUpperLayer(2, identity; use_bias = false) -A = UpperTriangular([1], 2) +A = StrictlyUpperTriangular([1], 2) @test l(ones(eltype(A), 2), (weight = A,)) == [2, 1] diff --git a/test/map_to_cpu_tests.jl b/test/map_to_cpu_tests.jl index c7d085421..0e9c1ccf5 100644 --- a/test/map_to_cpu_tests.jl +++ b/test/map_to_cpu_tests.jl @@ -12,7 +12,7 @@ using Random using Test import GeometricOptimizers: StiefelManifold, SymmetricMatrix, SkewSymMatrix, - LowerTriangular, UpperTriangular + StrictlyLowerTriangular, StrictlyUpperTriangular Random.seed!(1234) @@ -22,7 +22,7 @@ const N, n = 6, 3 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))) + L3 = (L = StrictlyLowerTriangular(randn(n, n)), U = StrictlyUpperTriangular(randn(n, n))) )) back = map_to_cpu(ps) @@ -35,8 +35,8 @@ const N, n = 6, 3 @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 + @test back.L3.L isa StrictlyLowerTriangular + @test back.L3.U isa StrictlyUpperTriangular # 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 diff --git a/test/optimizers/adam_with_learning_rate_decay.jl b/test/optimizers/adam_with_learning_rate_decay.jl index 12c590f94..145467b4c 100644 --- a/test/optimizers/adam_with_learning_rate_decay.jl +++ b/test/optimizers/adam_with_learning_rate_decay.jl @@ -3,6 +3,7 @@ using GeometricMachineLearning using GeometricMachineLearning: ResNetLayer, params using LinearAlgebra: I using Test +import GeometricOptimizers import Random Random.seed!(123) @@ -28,7 +29,7 @@ function train_network(; n_epochs = 2048) # pairing its own `Optimizer` takes, so it splats rather than being passed positionally. The # element type is positional and defaults to `Float64` here, where GML's own version took it # from `η₁` and so defaulted to `Float32`. - o₂ = Optimizer(nn₂; AdamOptimizerWithDecay(n_epochs, eltype(dl))...) + o₂ = Optimizer(nn₂; AdamOptimizerWithDecay(n_epochs)...) batch = Batch(5, 1) loss = FeedForwardLoss() @@ -59,7 +60,7 @@ function train_manifold_network(; n_epochs = 128) arch = Chain(StiefelLayer(1, 20), Dense(20, 20, tanh), Dense(20, 1, identity)) nn = NeuralNetwork(arch, CPU(), eltype(dl)) - o = Optimizer(nn; AdamOptimizerWithDecay(n_epochs, eltype(dl))...) + o = Optimizer(nn; AdamOptimizerWithDecay(n_epochs)...) loss_array = o(nn, dl, Batch(5, 1), n_epochs, FeedForwardLoss(); show_progress = false) @test all(isfinite, loss_array) @@ -72,32 +73,31 @@ end @doc raw""" The schedule is walked from ``t = 1``, not from ``t = 0``. -`optimization_step!` increments `opt.iterations` *before* it reads the step size, so the first step -of a run takes ``\alpha(1) = \gamma\eta_1`` and not ``\alpha(0) = \eta_1``. That is how the pre-0.5 -`AdamOptimizerWithDecay` counted — it incremented `o.step` before `update!` — and how +`optimization_step!` increments the iteration number *before* it reads the step size, so the first +step of a run takes ``\alpha(1) = \gamma\eta_1`` and not ``\alpha(0) = \eta_1``. That is how the +pre-0.5 `AdamOptimizerWithDecay` counted — it incremented `o.step` before `update!` — and how `DecayingStatic` counts, because `GeometricOptimizers.solve!` calls `increase_iteration_number!` before `solver_step!`. Reading before incrementing put every step of a run one place early in the schedule, which `test/adam_optimizer_with_decay.jl` upstream asserts does not happen. """ function schedule_starts_at_one(; n_epochs = 100, η₁ = 1e-2, η₂ = 1e-6) - method = AdamOptimizerWithDecay(n_epochs, Float64; η₁ = η₁, η₂ = η₂) + method = AdamOptimizerWithDecay(n_epochs; η₁ = η₁, η₂ = η₂) o = Optimizer(NetworkParameters((weight = zeros(2, 2),)); method...) γ = exp(log(η₂ / η₁) / n_epochs) - @test o.step_size isa DecayingStatic - @test o.iterations == 0 + @test o.training.linesearch isa DecayingStatic + @test GeometricOptimizers.iteration_number(o.training.state) == 0 for t in 1:4 - o.iterations += 1 - @test GeometricMachineLearning._current_step_size(o, o.iterations) ≈ η₁ * γ^t + @test GeometricOptimizers.step_size(o.training.linesearch, t) ≈ η₁ * γ^t end # and the same thing through the public entry point: one step of a Euclidean parameter with a # gradient of `1` moves it by `α₁ / (√1 + δ) ≈ α₁`, so the distance travelled reports the α used ps = NetworkParameters((weight = zeros(2, 2),)) opt = Optimizer(ps; method...) - optimization_step!(opt, GlobalSection(ps), ps, (weight = ones(2, 2),)) - @test opt.iterations == 1 + optimization_step!(ps, opt, (weight = ones(2, 2),)) + @test GeometricOptimizers.iteration_number(opt.training.state) == 1 @test abs(ps.weight[1, 1]) ≈ η₁ * γ rtol = 1e-6 end diff --git a/test/optimizers/composite_method.jl b/test/optimizers/composite_method.jl new file mode 100644 index 000000000..b58e8787e --- /dev/null +++ b/test/optimizers/composite_method.jl @@ -0,0 +1,55 @@ +using GeometricMachineLearning, Test, KernelAbstractions +import GeometricOptimizers +import Random + +# A network with a Stiefel layer beside a Euclidean one is the shape `CompositeMethod` exists for: +# `ScalarMomentAdam` steps a single `StiefelManifold`, so the dense layer needs another method. The +# per-leaf caches and states are `GeometricOptimizers`'; what is pinned here is that `Optimizer` +# hands a composite through to them unchanged. + +function stiefel_beside_dense(T) + Random.seed!(7) + NeuralNetwork(Chain(StiefelLayer(6, 3), Dense(6, 6, tanh)), KernelAbstractions.CPU(), T) +end + +function minibatch_gradient(T, k) + Random.seed!(100 + k) + (L1 = (weight = randn(T, 6, 3),), L2 = (W = randn(T, 6, 6), b = randn(T, 6))) +end + +@testset "a composite gives every leaf the state of the method it selects" begin + nn = stiefel_beside_dense(Float32) + opt = Optimizer(CompositeMethod(; manifold = ScalarMomentAdam(), array = Adam()), nn) + states = opt.training.state.states + @test opt.training.state isa CompositeState{Float32} + @test states.L1.weight isa ScalarMomentAdamState{Float32} + @test states.L2.W isa AdamState{Float32} + @test states.L2.b isa AdamState{Float32} + + ps = GeometricMachineLearning.params(nn) + for k in 1:3 + optimization_step!(ps, opt, minibatch_gradient(Float32, k)) + end + @test ps.L1.weight isa StiefelManifold + @test GeometricOptimizers.check(ps.L1.weight) < 1.0f-5 + @test GeometricOptimizers.iteration_number(opt.training.state) == 3 +end + +@testset "a composite of one method is that method, exactly" begin + # Each optimizer is built straight after its network, so both draw their section completions + # from the same position of the seed. + whole = stiefel_beside_dense(Float64) + whole_opt = Optimizer(Adam(), whole) + per_leaf = stiefel_beside_dense(Float64) + per_leaf_opt = Optimizer(CompositeMethod(; manifold = Adam(), array = Adam()), per_leaf) + for k in 1:5 + optimization_step!(GeometricMachineLearning.params(whole), whole_opt, + minibatch_gradient(Float64, k)) + optimization_step!(GeometricMachineLearning.params(per_leaf), per_leaf_opt, + minibatch_gradient(Float64, k)) + end + a, b = GeometricMachineLearning.params(whole), GeometricMachineLearning.params(per_leaf) + @test a.L1.weight.A == b.L1.weight.A + @test a.L2.W == b.L2.W + @test a.L2.b == b.L2.b +end diff --git a/test/optimizers/gml_gradient_dispatch.jl b/test/optimizers/gml_gradient_dispatch.jl deleted file mode 100644 index 85a9bcb6c..000000000 --- a/test/optimizers/gml_gradient_dispatch.jl +++ /dev/null @@ -1,129 +0,0 @@ -# What `_GMLGradient` is called on, and by which method. -# -# `_GMLGradient` is a `GeometricOptimizers.Gradient`, so its functor methods share a method table with -# upstream's. Two of them overlap on a shape this package actually produces, and the overlap is an -# *ambiguity* rather than a choice — which is a run-time error several frames into a solve, not a -# signature a reader would notice. So it is asserted here rather than left to whichever call path -# happens to reach it. - -using GeometricMachineLearning -using GeometricOptimizers: StiefelManifold, GrassmannManifold, rgrad -using NeuralNetworkParameters: NetworkParameters -using Test -import Random - -Random.seed!(1234) - -GML = GeometricMachineLearning - -# A bare `Manifold`. `_GMLGradient{T} <: GeometricOptimizers.Gradient{T}` and `Manifold{T} <: -# AbstractMatrix{T}`, so `(::_GMLGradient{T})(::AbstractArray{T})` and upstream's -# `(::Gradient{T})(::Manifold{T})` are ambiguous here: neither is more specific. Without a method of -# its own this call is `MethodError: ... is ambiguous`. -# -# The value is this package's answer and not upstream's. Upstream's body is -# `rgrad(x, reshape(grad(vec(x)), size(x)...))`, which evaluates an *inner* gradient — and a -# `_GMLGradient` has none, since the whole point of it is that the Euclidean gradient was computed -# already and is being carried in `dp`. -@testset "a bare Manifold reaches this package's method, not upstream's" begin - for M in (StiefelManifold, GrassmannManifold) - Y = rand(M, 4, 2) - dp = randn(4, 2) - g = GML._GMLGradient{Float64, typeof(dp)}(dp) - - @test g(Y) == rgrad(Y, dp) - @test which(g, Tuple{typeof(Y)}).module === GeometricMachineLearning - end -end - -# A bare `AbstractVector`. `_GMLGradient{T} <: GeometricOptimizers.Gradient{T}` and `SimpleSolvers` -# has `(::Gradient{T})(::AbstractVector{T})`, so without a method of its own, -# `(::_GMLGradient{T})(::AbstractArray{T})` and that one are equally specific on a `Vector` -- -# the call is `MethodError: ... is ambiguous`, same class of defect as the `Manifold` case above. -@testset "a bare AbstractVector reaches this package's method, not upstream's" begin - dp = randn(3) - g = GML._GMLGradient{Float64, typeof(dp)}(dp) - - @test g(randn(3)) === dp - @test which(g, Tuple{Vector{Float64}}).module === GeometricMachineLearning -end - -# The two shapes the optimizer actually hands it, so that the method above is not the only one -# covered and a later narrowing of either shows up here. -@testset "a wrapped layer and a plain array" begin - dp = randn(3, 3) - g = GML._GMLGradient{Float64, typeof(dp)}(dp) - # a plain array leaf: the gradient is already Euclidean, so it passes through - @test g(randn(3, 3)) === dp - - ps = NetworkParameters((weight = rand(StiefelManifold, 4, 2),)) - dps = (weight = randn(4, 2),) - gps = GML._GMLGradient{Float64, typeof(dps)}(dps) - out = gps(ps) - # rebuilt in the shape of its *first* argument -- a container, which is what - # `GeometricOptimizers._copyto!` has a method for - @test out isa NetworkParameters - @test out.weight == rgrad(ps.weight, dps.weight) -end - -# `_is_layer` is the rule for how a network is split into caches, so its boundary cases are stated -# rather than inherited. An empty set is *zero* caches' worth and descends; `all` over an empty -# collection would otherwise make it vacuously one. -@testset "_is_layer at the boundary" begin - @test GML._is_layer((W = randn(2, 2), b = randn(2))) - @test GML._is_layer(NetworkParameters((W = randn(2, 2), b = randn(2)))) - # a set whose values are branches is a subtree, and gets one cache per layer - @test !GML._is_layer((L1 = (W = randn(2, 2),), L2 = (W = randn(2, 2),))) - @test !GML._is_layer(NamedTuple()) - @test !GML._is_layer(NetworkParameters(NamedTuple())) - @test !GML._is_layer(randn(2, 2)) - - # a layer whose weights do not share an element type is still one cache, with `T` the promotion - @test GML._is_layer((W = randn(Float32, 2, 2), b = randn(Float64, 2))) -end - -# `_as_go_solution` is the wrap itself, and the property that makes the wrap free: it shares the leaf -# arrays, so the optimizer step writes through to the network's own weights and nothing is copied -# back. That is what makes this a boundary change rather than a change of behaviour, and it is the one -# claim of this release that a reader cannot check by reading a signature. -@testset "_as_go_solution wraps a layer and shares its arrays" begin - W, b = randn(2, 2), randn(2) - layer = (W = W, b = b) - - wrapped = GML._as_go_solution(layer) - @test wrapped isa NetworkParameters - @test wrapped.W === W # shared, not copied - @test wrapped.b === b - - # already in the right shape: passed through untouched - ps = NetworkParameters(layer) - @test GML._as_go_solution(ps) === ps - Y = rand(StiefelManifold, 4, 2) - @test GML._as_go_solution(Y) === Y - v = randn(3) - @test GML._as_go_solution(v) === v - - # a subtree is not a layer, so it is left alone for the tree branch to descend - tree = (L1 = (W = randn(2, 2),), L2 = (W = randn(2, 2),)) - @test GML._as_go_solution(tree) === tree -end - -# `_get_contents` unwraps whatever shape the pullback handed back. Six methods, because a reverse pass -# produces a container or the bare `NamedTuple` it wraps, and `Zygote` may hand either back bare, in a -# one-tuple, or in a one-element vector. Each is exercised here: splitting a union into a method per -# shape means each method needs its own case, or the split trades one tested path for six untested -# ones. -@testset "_get_contents unwraps every shape a pullback returns" begin - nt = (L1 = (W = randn(2, 2), b = randn(2)),) - ps = NetworkParameters(nt) - - for x in (nt, ps) - @test GML._get_contents(x) === x - @test GML._get_contents((x,)) === x - @test GML._get_contents([x]) === x - end - - # more than one set is a caller error and says so, rather than silently taking the first - @test_throws ArgumentError GML._get_contents([ps, ps]) - @test_throws ArgumentError GML._get_contents([nt, nt]) -end diff --git a/test/optimizers/gradient_optimizer.jl b/test/optimizers/gradient_optimizer.jl index 376e946bd..12d68225c 100644 --- a/test/optimizers/gradient_optimizer.jl +++ b/test/optimizers/gradient_optimizer.jl @@ -18,8 +18,7 @@ function gradient_optimizer_euclidean(N; n_steps = 20, step_size = 1e-3) o = Optimizer(GradientMethod(), ps; step_size = step_size) for _ in 1:n_steps ∇L = Zygote.gradient(loss, ps)[1] - λY = GlobalSection(ps) - optimization_step!(o, λY, ps, ∇L) + optimization_step!(ps, o, ∇L) end loss2 = loss(ps) @test loss1 > loss2 @@ -38,8 +37,7 @@ function gradient_optimizer_stiefel(N, n; n_steps = 20, step_size = 1e-3) o = Optimizer(GradientMethod(), ps; step_size = step_size) for _ in 1:n_steps ∇L = Zygote.gradient(loss, ps)[1] - λY = GlobalSection(ps) - optimization_step!(o, λY, ps, ∇L) + optimization_step!(ps, o, ∇L) end loss2 = loss(ps) @test loss1 > loss2 diff --git a/test/optimizers/momentum_optimizer.jl b/test/optimizers/momentum_optimizer.jl index 0351140a1..9359c806f 100644 --- a/test/optimizers/momentum_optimizer.jl +++ b/test/optimizers/momentum_optimizer.jl @@ -16,11 +16,10 @@ growing one that never saturates, which is what this test rules out. """ function momentum_matches_geometric_sum(; α = 0.9, η = 0.1, n_steps = 10) ps = NetworkParameters((A = zeros(2, 2),)) - o = Optimizer(MomentumMethod(α), ps; step_size = η) - λY = GlobalSection(ps) + o = Optimizer(MomentumMethod(; α = α), ps; step_size = η) previous = 0.0 for t in 1:n_steps - optimization_step!(o, λY, ps, (A = ones(2, 2),)) + optimization_step!(ps, o, (A = ones(2, 2),)) step = previous - ps.A[1, 1] previous = ps.A[1, 1] @test step ≈ η * (1 - α^t) / (1 - α) @@ -38,10 +37,9 @@ function momentum_optimizer_stiefel(N, n; n_steps = 50, step_size = 1e-3, α = 0 loss(ps) = norm(ps.Y * ps.Y' - B)^2 ps = NetworkParameters((Y = rand(StiefelManifold, N, n),)) loss1 = loss(ps) - o = Optimizer(MomentumMethod(α), ps; step_size = step_size) - λY = GlobalSection(ps) + o = Optimizer(MomentumMethod(; α = α), ps; step_size = step_size) for _ in 1:n_steps - optimization_step!(o, λY, ps, Zygote.gradient(loss, ps)[1]) + optimization_step!(ps, o, Zygote.gradient(loss, ps)[1]) end @test loss(ps) < loss1 @test ps.Y' * ps.Y ≈ I diff --git a/test/optimizers/optimization_step.jl b/test/optimizers/optimization_step.jl index 781d625a0..b38a068bf 100644 --- a/test/optimizers/optimization_step.jl +++ b/test/optimizers/optimization_step.jl @@ -17,8 +17,7 @@ function optimization_step_test(N, n, T) o = Optimizer(m, ps) ps2 = deepcopy(ps) - λY = GlobalSection(ps) - optimization_step!(o, λY, ps, dx) + optimization_step!(ps, o, dx) @test typeof(ps[1].weight) <: StiefelManifold for (layers1, layers2) in zip(values(ps), values(ps2)) for key in keys(layers1) @@ -35,23 +34,16 @@ for N in 4:N_max end 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 +# The whole network is one `GeometricOptimizers.TrainingOptimizer`: one cache and one state over the +# whole parameter set, and not one per layer as before 0.9. +@testset "one cache and one state over the whole network" 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 + @test o.training isa GeometricOptimizers.TrainingOptimizer + @test o.training.cache isa GeometricOptimizers.AdamCache{Float32} + @test o.training.state isa AdamState{Float32} + @test keys(GeometricOptimizers.solution(o.training.state)) == + keys(GeometricMachineLearning.params(nn)) end diff --git a/test/optimizers/psd_optim.jl b/test/optimizers/psd_optim.jl index 8afad17a7..3ccc27d48 100644 --- a/test/optimizers/psd_optim.jl +++ b/test/optimizers/psd_optim.jl @@ -26,7 +26,7 @@ A = [0.06476993260924702 0.8369280855305259 0.6245358125914054 0.140729967064923 """ This tests if the optimizers can find the optimal PSD solution. """ -function svd_test(A, n, train_steps = 1000, tol = 1e-1; retraction = cayley) +function svd_test(A, n, train_steps = 1000, tol = 1e-1; retraction = Cayley()) # Seeded here and not once at the top of the file, for the reason given at the same place in # `svd_optim.jl`: `svd_test` is called twice, and with a single top-level seed the second call # starts from whatever RNG state the first one happened to leave behind, which makes the @@ -69,13 +69,12 @@ function train_network!(o::Optimizer, model::Chain, ps::NetworkParameters, for _ in 1:train_steps dx = Zygote.gradient(error, ps)[1] - λY = GlobalSection(ps) - optimization_step!(o, λY, ps, dx) + optimization_step!(ps, o, dx) #println(error(ps)) end ps[1].weight, ps[2].weight, error(ps) end -for retraction in (geodesic, cayley) +for retraction in (Geodesic(), Cayley()) svd_test(A, 4, retraction = retraction) end diff --git a/test/optimizers/step_size_element_type.jl b/test/optimizers/step_size_element_type.jl deleted file mode 100644 index a2a13658a..000000000 --- a/test/optimizers/step_size_element_type.jl +++ /dev/null @@ -1,52 +0,0 @@ -# `_leaf_optim_step!` scales `direction(cache)` by `T(step_size)`, matching the three -# `_euclidean_update!` methods a few lines below it. The funnel (`_current_step_size`) always -# hands a `Float64`, no matter what `T` the parameters are, so converting first is what keeps a -# `Float32` layer scaled at `Float32` precision instead of `Float64` precision rounded back on -# write. See `CHANGELOG.md` for the fix this guards. -# -# The test is invariance rather than a fixed expected value: the leaf step must give the *same* -# result whether the funnel hands it a `Float64` step size or one already rounded to `T`. Only -# `T = Float32` can tell the two calls apart -- `T(step_size)` is a no-op when `T` is already -# `Float64`, so a `Float64` combination holds whatever the conversion does and checks nothing. -# `Float64` is therefore not looped over here. -# -# The `seed = 2` and the `step_size = 0.1` are chosen rather than arbitrary. Whether a `Float64`- -# and a `Float32`-scaled direction round to different bit patterns depends on the actual values, -# so for most seeds the two calls agree however the step size is scaled, and the test then passes -# without discriminating. With these values, and with the conversion dropped, both `Adam` and -# `MomentumMethod` separate by `2.9802322f-8` -- one `Float32` ULP at this magnitude -- and each -# method is checked on its own. - -using GeometricMachineLearning -using GeometricOptimizers -using Test -import Random - -GML = GeometricMachineLearning - -function test_leaf_step_matches_prerounded_step_size(method, seed::Int, step_size::Float64) - T = Float32 - Random.seed!(seed) - ps_a = (weight = rand(T, 6, 6),) - dp = (weight = rand(T, 6, 6),) - ps_b = deepcopy(ps_a) - λY_a = GML.GlobalSection(ps_a) - λY_b = GML.GlobalSection(ps_b) - - cache_a = GML._make_optimizer_cache(method, ps_a) - state_a = GML._make_optimizer_state(method, ps_a) - cache_b = GML._make_optimizer_cache(method, ps_b) - state_b = GML._make_optimizer_state(method, ps_b) - - # the funnel's usual `Float64` step size ... - GML._leaf_optim_step!( - cache_a, state_a, dp, ps_a, λY_a, method, GeometricOptimizers.cayley, step_size) - # ... versus the same value pre-rounded to `T`. - GML._leaf_optim_step!( - cache_b, state_b, dp, ps_b, λY_b, method, GeometricOptimizers.cayley, T(step_size)) - - @test ps_a.weight == ps_b.weight -end - -test_leaf_step_matches_prerounded_step_size(GeometricOptimizers.Adam(), 2, 0.1) -test_leaf_step_matches_prerounded_step_size(GeometricOptimizers.MomentumMethod(0.5), 2, 0.1) diff --git a/test/optimizers/structured_array_parameters.jl b/test/optimizers/structured_array_parameters.jl index 4ce1a08d3..c9a000d2c 100644 --- a/test/optimizers/structured_array_parameters.jl +++ b/test/optimizers/structured_array_parameters.jl @@ -8,7 +8,8 @@ Random.seed!(123) @doc raw""" Every optimizer method has to work with the structured matrix types GML uses as *ordinary* weights: `SymmetricMatrix` (SympNet and symplectic attention layers), `SkewSymMatrix` (volume-preserving -attention) and `LowerTriangular`/`UpperTriangular` (volume-preserving feedforward layers). +attention) and `StrictlyLowerTriangular`/`StrictlyUpperTriangular` (volume-preserving feedforward +layers). These have their own storage and no `setindex!`, so they need a type-preserving `similar` and elementwise `_add!`/`_rac!`/`_square!`/`_div!`/`_rmul!` written on the free parameters; without them @@ -31,7 +32,7 @@ function optimizer_runs(architecture, batch, input_dim; T = Float64, n_epochs = # `AdamOptimizerWithDecay` is a `(algorithm, linesearch)` pairing rather than a method, so it # goes in through the keyword constructor nn, dl = nn_and_dl() - o = Optimizer(nn; AdamOptimizerWithDecay(n_epochs, T)...) + o = Optimizer(nn; AdamOptimizerWithDecay(n_epochs)...) loss_array = o(nn, dl, batch, n_epochs; show_progress = false) @test length(loss_array) == n_epochs @test all(isfinite, loss_array) @@ -41,7 +42,7 @@ end # then fails to match them against the parameter @testset "structured weights keep their type in `similar`" begin for A in (rand(SymmetricMatrix{Float64}, 4), rand(SkewSymMatrix{Float64}, 4), - rand(LowerTriangular{Float64}, 4), rand(UpperTriangular{Float64}, 4)) + rand(StrictlyLowerTriangular{Float64}, 4), rand(StrictlyUpperTriangular{Float64}, 4)) @test typeof(similar(A)) == typeof(A) @test typeof(zero(A)) == typeof(A) @test size(similar(A)) == size(A) @@ -51,7 +52,7 @@ end @testset "`LASympNet` (SymmetricMatrix)" begin optimizer_runs(LASympNet(4), Batch(5), 4) end -@testset "`VolumePreservingFeedForward` (Lower/UpperTriangular)" begin +@testset "`VolumePreservingFeedForward` (the triangular matrices)" begin optimizer_runs(VolumePreservingFeedForward(4), Batch(5), 4) end @testset "`VolumePreservingTransformer` (SkewSymMatrix + triangular)" begin diff --git a/test/optimizers/svd_optim.jl b/test/optimizers/svd_optim.jl index 033819dee..47b7b24ed 100644 --- a/test/optimizers/svd_optim.jl +++ b/test/optimizers/svd_optim.jl @@ -25,7 +25,7 @@ A = [0.06476993260924702 0.8369280855305259 0.6245358125914054 0.140729967064923 Random.seed!(1234) -function svd_test(A, n, train_steps = 1000, tol = 1e-1; retraction = cayley) +function svd_test(A, n, train_steps = 1000, tol = 1e-1; retraction = Cayley()) # Seeded here and not once at the top of the file. `svd_test` is called twice, and with a single # top-level seed the second call starts from whatever RNG state the first one happened to leave # behind -- so the convergence assertions below turned on how much randomness the optimizer @@ -65,12 +65,11 @@ function train_network!(o::Optimizer, model::Chain, ps::NetworkParameters, for _ in 1:train_steps dx = Zygote.gradient(error, ps)[1] - λY = GlobalSection(ps) - optimization_step!(o, λY, ps, dx) + optimization_step!(ps, o, dx) end ps.L1.weight, ps.L2.weight, error(ps) end -for retraction in (geodesic, cayley) +for retraction in (Geodesic(), Cayley()) svd_test(A, 3, retraction = retraction) end diff --git a/test/quality/aqua.jl b/test/quality/aqua.jl index 8ce09ad3d..00ad00956 100644 --- a/test/quality/aqua.jl +++ b/test/quality/aqua.jl @@ -3,43 +3,20 @@ using ExplicitImports using GeometricMachineLearning using Test -# Aqua's package-level checks. Six of the eight `Aqua.test_all` enables by default run; `ambiguities` -# and `piracies` are switched off by name, and the reason is below rather than in the changelog -# alone, because switching one back on turns the suite red on the spot. +# Aqua's package-level checks, all eight that `Aqua.test_all` enables by default. # # `Aqua.test_all` wraps each of its checks in a `@testset` of its own and adds no enclosing one, and # a testset with no parent finalises as soon as it closes -- so called bare, the first failing check # throws a `TestSetException` and the rest never run. On the suite path the `@safetestset` in # `runtests.jl` is that parent. The `@testset` below is that parent when this file is run on its -# own, and it is also what lets the two counts at the end be further assertions rather than -# top-level `@test`s that abort the file the moment one fails. +# own. # -# WHAT IS SWITCHED OFF, AND WHY IT IS NOT A `broken = true`. -# -# `piracies` reports 3 methods, and all 3 are genuine under Aqua's definition -- the function and -# every argument type belong to other modules. They are not a count waiting to be triaged. All -# three are the same thing: a three-argument functor for a *layer* type that this package does not -# own, `AbstractNeuralNetworks`' `Dense` and `Linear`, added so that the layer accepts a 3-tensor -# input as well as a matrix. -# -# Dense{M, N, true} on an AbstractArray{T, 3} src/layers/resnet.jl:63 -# Dense{M, N, false} on an AbstractArray{T, 3} src/layers/resnet.jl:67 -# Linear{M, N} on an AbstractArray{T, 3} src/layers/resnet.jl:71 -# -# Closing these means either upstream gaining the 3-tensor methods or this package wrapping the two -# layer types, and both are an API change rather than a tidy-up. They are *B7* under -# `## Open Issues` in `CHANGELOG.md` with that reasoning. -# -# `ambiguities` reports 1, and the same 1 whether this package is loaded alone or with everything -# `runtests.jl` loads before this file. It is a `Dense{M, N, true}` functor against -# `AbstractNeuralNetworks.Affine`'s, and it is benign: `Dense` is not a subtype of `Affine` and the -# two have no common instance, so no call can reach the pair. It is the same `resnet.jl:63` method -# as the first piracy above. -# -# The two settings agreeing is what makes the gate below trustworthy, and it holds because no -# method of this package claims an argument position wholesale on `PoissonTensor`. That type is an -# `AbstractMatrix{T}`, so such a method meets every special-array method another package writes -# against `AbstractMatrix`: +# `ambiguities` and `piracies` were switched off until 0.9, with a count gate for each in their +# place: three pirated 3-tensor functors on `AbstractNeuralNetworks`' `Dense` and `Linear`, and one +# ambiguity between the first of them and `Affine`'s. `AbstractNeuralNetworks` 0.9 defines those +# methods itself, so they are deleted here and both checks run. No method of this package claims an +# argument position wholesale on `PoissonTensor`, which is an `AbstractMatrix{T}` and would +# otherwise meet every special-array method another package writes against `AbstractMatrix`: # # `*(𝕁::PoissonTensor{T}, v::Strided…)` takes a `Strided…` right-hand side, which excludes each # special array type `X` that ArrayLayouts, FillArrays, Symbolics and GeometricOptimizers @@ -48,29 +25,8 @@ using Test # `BlockIndex` and `BandRangeType` methods BandedMatrices and BlockArrays add to # `AbstractMatrix`. Neither of those two loads with this package alone, so an untyped index # pair here is a collision visible only from inside the suite. -# -# `ambiguities` still stays off, because 1 is not 0 and Aqua's check has no way to exempt a pair. -# The assertion below replaces it and is strictly better here: it names the number, so it fails -# both when an ambiguity is added and when this one is closed without the comment going with it. -# -# Marking either check `broken = true` would leave a check that reports success while the defect -# stands, which is the failure mode this test suite's guards exist to remove. @testset "Aqua" begin - Aqua.test_all(GeometricMachineLearning; ambiguities = false, piracies = false) - - # A switched-off check detects nothing, so the count it would have reported drifts unobserved, - # and so does every `src` line named above. This is the gate. It fails when a piracy is added, - # and it fails when one is removed without the entry above and in `CHANGELOG.md` going with it. - @test length(Aqua.Piracy.hunt(GeometricMachineLearning)) == 3 - - # The ambiguity count gets the same gate, and it can carry one because the single pair is - # between a method of this package and one of `AbstractNeuralNetworks`'. The count is the same - # in isolation and in the suite, and nothing outside this repository's own `[compat]` moves it - # -- so a red here is a regression, not an unrelated upstream upgrade and not a reordering of - # this suite. A count that moved with ArrayLayouts', FillArrays', Symbolics', - # GeometricOptimizers', BandedMatrices' and BlockArrays' versions, and with *which* of them a - # given process had loaded, could not tell those apart, and could not be asserted. - @test length(Test.detect_ambiguities(GeometricMachineLearning; recursive = true)) == 1 + Aqua.test_all(GeometricMachineLearning) end # `ExplicitImports` answers a question Aqua does not ask, and the reason to gate on it is not diff --git a/test/runtests.jl b/test/runtests.jl index 9d1e773bc..5df2e8f0b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -59,8 +59,7 @@ if "core" in GROUPS @safetestset "Gradient optimizer tests" include("optimizers/gradient_optimizer.jl") @safetestset "Momentum optimizer tests" include("optimizers/momentum_optimizer.jl") @safetestset "Optimizers with structured (non-manifold) weights" include("optimizers/structured_array_parameters.jl") - @safetestset "_GMLGradient dispatch" include("optimizers/gml_gradient_dispatch.jl") - @safetestset "The GO-native leaf step scales in the parameter's own element type" include("optimizers/step_size_element_type.jl") + @safetestset "CompositeMethod, one method per leaf" include("optimizers/composite_method.jl") @safetestset "_custom_mul: NetworkParameters gradient structure" include("layers/double_multiplication_network_parameters_gradient.jl") @safetestset "Symplectic attention: NetworkParameters gradient structure" include("layers/symplectic_attention_network_parameters_gradient.jl") @safetestset "map_to_cpu" include("map_to_cpu_tests.jl")