From 0bd71a457e915bca895e4d5c6258273bc20b518a Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Sun, 25 Oct 2026 18:54:02 +0100 Subject: [PATCH 1/7] Add apply_parametric, ParametricLoss and ParametricDataLoader Port the parametric core of pull request #207 onto main: the concatenation of the system parameters with the state, apply_parametric, ParametricLoss with its two ZygotePullback methods, and ParametricDataLoader on main's gather, batch_over_two_axes and train!. The three names are public and not exported. train! on a ParametricDataLoader trains on ParametricLoss() without a loss. Relates to #207. Co-authored-by: benedict-96 Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 + docs/src/data_loader/data_loader.md | 1 + docs/src/reduced_order_modeling/losses.md | 2 + src/GeometricMachineLearning.jl | 8 ++ src/data_loader/parametric_data_loader.jl | 102 +++++++++++++ src/data_loader/train.jl | 21 ++- src/loss/losses.jl | 26 ++++ src/parametric.jl | 110 ++++++++++++++ src/pullbacks/zygote_pullback.jl | 22 ++- test/data_loader/parametric_data_loader.jl | 158 +++++++++++++++++++++ test/helpers/parametric_affine.jl | 35 +++++ test/loss/parametric_loss.jl | 82 +++++++++++ test/parametric.jl | 94 ++++++++++++ test/runtests.jl | 3 + 14 files changed, 657 insertions(+), 9 deletions(-) create mode 100644 src/data_loader/parametric_data_loader.jl create mode 100644 src/parametric.jl create mode 100644 test/data_loader/parametric_data_loader.jl create mode 100644 test/helpers/parametric_affine.jl create mode 100644 test/loss/parametric_loss.jl create mode 100644 test/parametric.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index ba2999f81..2956eef92 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -99,6 +99,8 @@ breaking release). ### Added +- **`apply_parametric`, `ParametricLoss` and `ParametricDataLoader` train a network whose forward pass takes the parameters of the system beside its input.** The three names are `public` and not exported, so a caller qualifies them or names them in a `using GeometricMachineLearning: …`. `apply_parametric(model, x, μ, ps)` applies a `Chain` to one sample — a vector or a `(q, p)` pair of vectors — with one `μ`, or to a batch — a `2d × 1 × k` tensor or a `(q, p)` pair of `d × 1 × k` tensors — with a vector of `k` values of `μ`, sample by sample; the output has the shape of the input. The `Chain` receives the pair `(x, μ)`, and the internal `concatenate_array_with_parameters` puts the flattened `μ` below the state in the element type of the state, side by side for a matrix and a vector of `μ`, and repeated in every column for a matrix and one `μ`. `ParametricLoss()(model, ps, input, output, μ)` is the relative norm of the error of `apply_parametric`, and `ZygotePullback` takes the triple `(input, output, μ)` with one `μ` or a vector of them. `ParametricDataLoader(data, parameters)` takes a `2d × nt × m` array or a `(q, p)` pair of `d × nt × m` arrays with a vector of `m` named tuples, and `ParametricDataLoader(ensemble_solution)` takes an `EnsembleSolution` of `(q, p)` states with the parameters of its problem. A `Batch(b)` draws each pair of a start position and a trajectory once per epoch, as for `DataLoader`, and a minibatch is the state, the state one time step later and the `μ` of each sample. `train!` takes a `ParametricDataLoader` and trains on `ParametricLoss()` when no `loss` is given. A `Batch` with a sequence length other than one, a parameter vector whose length is not the number of trajectories, no trajectory, too few time steps for the batch, and a `μ` whose element type is not the data's raise an `ArgumentError`. The code is ported from pull request #207 by Benedikt Brantner (relates to #207); there `concatenate_array_with_parameters` promoted a `Float32` state to `Float64`, rejected a matrix of more than one column with one `μ`, and the loader and the batch checked their counts with `@assert`. + - **`KernelAbstractions.get_backend` is defined for `PoissonTensor`.** GeometricOptimizers 0.9 asks every operand of a checked product for its backend, so `𝕁 * Y` with a `StiefelManifold` raised `ArgumentError: Implement KernelAbstractions.get_backend(::PoissonTensor...)` without it. - **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, in the testset "_mul and tensor_tensor_mul on Metal match the host", `_mul` with a dense and a `SymmetricMatrix` weight from the left, a dense one from the right, and `tensor_tensor_mul` 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 and Julia 1.13.1, both testsets pass in a cold process: 4 of 4 assertions in the `PoissonTensor` testset and 8 of 8 in the "_mul and tensor_tensor_mul on Metal match the host" testset. diff --git a/docs/src/data_loader/data_loader.md b/docs/src/data_loader/data_loader.md index 2cb8a7bd7..340f4fb6d 100644 --- a/docs/src/data_loader/data_loader.md +++ b/docs/src/data_loader/data_loader.md @@ -165,4 +165,5 @@ Batch(::Int, ::Int, ::Int) GeometricMachineLearning.number_of_batches GeometricMachineLearning.convert_input_and_batch_indices_to_array GeometricMachineLearning.accuracy +GeometricMachineLearning.ParametricDataLoader ``` diff --git a/docs/src/reduced_order_modeling/losses.md b/docs/src/reduced_order_modeling/losses.md index 256c4673c..10d0d15e1 100644 --- a/docs/src/reduced_order_modeling/losses.md +++ b/docs/src/reduced_order_modeling/losses.md @@ -46,6 +46,8 @@ where ``\mathbf{x}^{(t)}`` is the solution of the FOM at point ``t`` and ``\math TransformerLoss AutoEncoderLoss ReducedLoss +GeometricMachineLearning.ParametricLoss +GeometricMachineLearning.apply_parametric projection_error reduction_error ``` diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index 747610220..6390dd02c 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -11,6 +11,9 @@ using AbstractNeuralNetworks # `NamedTuple` of layers is walked with `Base.map`, which needs nothing from anybody. import NeuralNetworkParameters: NetworkParameters using NeuralNetworkParameters: mapstorage +# `flatten` puts the parameters of a system into the input of a network that takes them +# (`src/parametric.jl`), and `parameter_eltype` checks their element type against the data's. +using NeuralNetworkParameters: flatten, parameter_eltype using ChainRulesCore using GeometricBase using GeometricSolutions: GeometricSolution, EnsembleSolution, TimeSeries @@ -205,6 +208,7 @@ export solve!, encoder, decoder export iterate +include("parametric.jl") include("loss/losses.jl") export AbstractPullback, ZygotePullback @@ -214,8 +218,12 @@ export DataLoader export Batch, train! include("data_loader/tensor_assign.jl") include("data_loader/batch.jl") +include("data_loader/parametric_data_loader.jl") include("data_loader/train.jl") +# The training of a network that takes the parameters of the system beside its input. +public apply_parametric, ParametricLoss, ParametricDataLoader + include("reduced_system/reduced_system.jl") export HRedSys, reduction_error, projection_error, integrate_reduced_system, diff --git a/src/data_loader/parametric_data_loader.jl b/src/data_loader/parametric_data_loader.jl new file mode 100644 index 000000000..ea51f6eb8 --- /dev/null +++ b/src/data_loader/parametric_data_loader.jl @@ -0,0 +1,102 @@ +@doc raw""" + ParametricDataLoader(data, parameters) + ParametricDataLoader(ensemble_solution) + +Make a loader of trajectories of a system together with the parameters ``\mu`` of the system that +each trajectory belongs to. + +`data` is a `2d × nt × m` array, or a `(q, p)` named tuple of two `d × nt × m` arrays: `m` +trajectories of `nt` time steps, with `q` above `p` where they are stacked. `parameters` is a vector +of `m` named tuples, the ``\mu`` of each trajectory. The data and every ``\mu`` have one element type. +An `EnsembleSolution` of a problem with `(q, p)` states gives its `q` and `p` and the parameters of its +problem. + +# The batch + +A `Batch(batch_size)` draws, in one epoch, each pair of a start position and a trajectory once, in +minibatches of `batch_size`; the last minibatch is shorter. A drawn minibatch is the triple +`(input, output, μs)`: the state at the start position, the state one time step later, and the +``\mu`` of each sample. [`train!`](@ref) on a `ParametricDataLoader` trains on +[`ParametricLoss`](@ref)`()` when no `loss` is given, so the network is applied with +[`apply_parametric`](@ref), one time step per sample. A `Batch` with a sequence length other than one +is an `ArgumentError`. + +A problem without system parameters is trained with a [`DataLoader`](@ref). +""" +struct ParametricDataLoader{ + T, AT <: ArrayOrNamedTuple{T}, PT <: AbstractVector{<:NamedTuple}} + input::AT + parameters::PT + input_dim::Int + input_time_steps::Int + n_params::Int + + function ParametricDataLoader(input::AT, + parameters::PT) where {T, + AT <: Union{AbstractArray{T, 3}, + NamedTuple{(:q, :p), <:Tuple{AbstractArray{T, 3}, AbstractArray{T, 3}}}}, + PT <: AbstractVector{<:NamedTuple}} + input_dim, input_time_steps, n_params = data_size(input) + n_params ≥ 1 && length(parameters) == n_params || throw(ArgumentError( + "the data hold $(n_params) trajectories and there are $(length(parameters)) " * + "system parameters; a `ParametricDataLoader` takes one for each trajectory, " * + "and at least one trajectory.")) + for μ in parameters + S = parameter_eltype(μ) + S === T || S === Union{} || + throw(ArgumentError( + "the data have the element type $(T) and a system parameter has the " * + "element type $(S); a `ParametricDataLoader` takes one element type.")) + end + new{T, AT, PT}(input, parameters, input_dim, input_time_steps, n_params) + end +end + +data_size(input::AbstractArray{<:Any, 3}) = size(input) + +function data_size(input::NamedTuple{(:q, :p)}) + size(input.q) == size(input.p) || throw(ArgumentError( + "`q` has the size $(size(input.q)) and `p` has the size $(size(input.p)).")) + (2 * size(input.q, 1), size(input.q, 2), size(input.q, 3)) +end + +function ParametricDataLoader(ensemble_solution::EnsembleSolution{T, + T1, + Vector{ST}}) where { + T, T1, TuT, TT <: TimeSeries{T1}, + ST <: GeometricSolution{T, T1, TT, NamedTuple{(:t, :q, :p, :q̇, :ṗ), TuT}}} + ParametricDataLoader( + (q = data_tensor([s.q for s in ensemble_solution.s]), + p = data_tensor([s.p for s in ensemble_solution.s])), + ensemble_solution.problem.parameters) +end + +function networkbackend(dl::ParametricDataLoader{T, <:AbstractArray}) where {T} + networkbackend(dl.input) +end +function networkbackend(dl::ParametricDataLoader{T, <:NamedTuple}) where {T} + networkbackend(dl.input.q) +end + +# The start positions are those with a target `prediction_window` steps later; the network takes one +# time step per sample, so the sequence length is one. +function (batch::Batch)(dl::ParametricDataLoader) + batch.seq_length == 1 || throw(ArgumentError( + "a `ParametricDataLoader` draws one time step per sample, and the `Batch` asks for a " * + "sequence length of $(batch.seq_length).")) + n = dl.input_time_steps - batch.prediction_window + n ≥ 1 || throw(ArgumentError( + "the `ParametricDataLoader` has $(dl.input_time_steps) time steps and the `Batch` asks " * + "for a prediction window of $(batch.prediction_window), which leaves no minibatch.")) + batch_over_two_axes(batch, n, dl.n_params) +end + +# The input at the start position, the target `seq_length` steps later, and the `μ` of the trajectory +# of each sample, which row 2 of the batch indices names. +function convert_input_and_batch_indices_to_array(dl::ParametricDataLoader, batch::Batch, + batch_indices_tuple::Vector{Tuple{Int, Int}}) + indices = convert_vector_of_tuples_to_matrix(networkbackend(dl), batch_indices_tuple) + (gather(dl.input, indices, 0, batch.seq_length), + gather(dl.input, indices, batch.seq_length, batch.prediction_window), + [dl.parameters[trajectory] for (_, trajectory) in batch_indices_tuple]) +end diff --git a/src/data_loader/train.jl b/src/data_loader/train.jl index 2ced7f23a..b3583eef7 100644 --- a/src/data_loader/train.jl +++ b/src/data_loader/train.jl @@ -2,7 +2,8 @@ train!(nn, opt, dl, batch, n_epochs; loss, show_progress = true) Train the neural network `nn` with the `GeometricOptimizers.TrainingOptimizer` `opt` on the data of -the [`DataLoader`](@ref) `dl` for `n_epochs` epochs, and return the loss history. +the [`DataLoader`](@ref) or [`ParametricDataLoader`](@ref) `dl` for `n_epochs` epochs, and return +the loss history. Each epoch draws the minibatches of `dl` with the [`Batch`](@ref) `batch` and takes one step per minibatch: `loss` computes the loss on the minibatch and its gradient, and @@ -15,8 +16,9 @@ gradient, a zero leaf included. Entry `i` of the history is the mean of the mini trained through a [`ZygotePullback`](@ref). Without it, a [`TransformerIntegrator`](@ref) with a `Batch{:Transformer}` trains on [`TransformerLoss`](@ref)`(batch)`, a [`NeuralNetworkIntegrator`](@ref) with a `Batch{:FeedForward}` on `FeedForwardLoss()`, and an -autoencoder with a `Batch{:FeedForward}` on [`AutoEncoderLoss`](@ref)`()`; any other combination -needs `loss`, and `train!` raises an `ArgumentError` without it. +autoencoder with a `Batch{:FeedForward}` on [`AutoEncoderLoss`](@ref)`()`, and any network on a +`ParametricDataLoader` on [`ParametricLoss`](@ref)`()`; any other combination needs `loss`, and +`train!` raises an `ArgumentError` without it. A network of HamiltonianNeuralNetworks.jl trains on vector-field data with `loss = SymbolicPullback(arch)` of that package, and on consecutive states with a loss such as its @@ -29,9 +31,9 @@ with the same `opt` resumes the training where the first one stopped; the only r `train!` draws are those of the minibatch shuffle. `n_epochs < 1` is an `ArgumentError`, raised before any step. """ -function train!(nn::NeuralNetwork, opt::TrainingOptimizer, dl::DataLoader, batch::Batch, - n_epochs::Integer; - loss::Union{NetworkLoss, AbstractPullback} = NetworkLoss(nn.architecture, batch), +function train!(nn::NeuralNetwork, opt::TrainingOptimizer, + dl::Union{DataLoader, ParametricDataLoader}, batch::Batch, n_epochs::Integer; + loss::Union{NetworkLoss, AbstractPullback} = default_loss(nn, dl, batch), show_progress::Bool = true) n_epochs ≥ 1 || throw(ArgumentError("`train!` needs at least one epoch; `n_epochs` is $(n_epochs).")) @@ -66,6 +68,13 @@ end _as_pullback(l::NetworkLoss) = ZygotePullback(l) _as_pullback(p::AbstractPullback) = p +# A `DataLoader` takes the default loss of the architecture and the batch; a `ParametricDataLoader` +# draws the system parameters of each sample, which `ParametricLoss` hands to the network. +function default_loss(nn::NeuralNetwork, ::DataLoader, batch::Batch) + NetworkLoss(nn.architecture, batch) +end +default_loss(::NeuralNetwork, ::ParametricDataLoader, ::Batch) = ParametricLoss() + # The default loss of an architecture and a batch. The two-argument methods live here and not beside # `NetworkLoss(::AutoEncoder)` in `src/loss/losses.jl`, because `Batch` is included after `src/loss/`. function AbstractNeuralNetworks.NetworkLoss( diff --git a/src/loss/losses.jl b/src/loss/losses.jl index 28ae38893..0820e5d28 100644 --- a/src/loss/losses.jl +++ b/src/loss/losses.jl @@ -226,3 +226,29 @@ function (loss::ReducedLoss)(model::Chain, ps::NetworkParameters, input::ArrayOrNamedTuple, output::ArrayOrNamedTuple) _compute_loss(loss.decoder(model(loss.encoder(input), ps)), output) end + +@doc raw""" + ParametricLoss() + +Make an instance of the loss of a network whose forward pass takes the parameters ``\mu`` of the +system beside its input. + +`ParametricLoss()(model, ps, input, output, μ)` applies the network with +[`apply_parametric`](@ref) and compares the result with `output` in the relative ``L_2`` norm: + +```math +\mathrm{loss}(\mathcal{NN}, \mathrm{input}, \mathrm{output}, \mu) = ||\mathcal{NN}(\mathrm{input}, \mu) - \mathrm{output}|| / ||\mathrm{output}||. +``` + +This is `FeedForwardLoss` with the system parameters passed to the network. `μ` is one +`OptionalParameters` for one sample, or a vector of them, one per sample of a batch, as a +[`ParametricDataLoader`](@ref) draws it. `train!` on a `ParametricDataLoader` trains on this loss +when no `loss` is given. +""" +struct ParametricLoss <: NetworkLoss end + +function (loss::ParametricLoss)(model::Chain, ps::NetworkParameters, + input::ArrayOrNamedTuple, output::ArrayOrNamedTuple, + μ::Union{OptionalParameters, AbstractVector}) + _compute_loss(apply_parametric(model, input, μ, ps), output) +end diff --git a/src/parametric.jl b/src/parametric.jl new file mode 100644 index 000000000..d11fd7510 --- /dev/null +++ b/src/parametric.jl @@ -0,0 +1,110 @@ +# The forward pass of a network that takes the parameters of the *system* — not of the network — +# beside its input. The system parameters `μ` are an `OptionalParameters`: a `NamedTuple`, or +# `NullParameters()` for a system without any. + +# The flat vector of the system parameters, in the element type `T` where it is given, together with +# the `ParameterLayout` that puts such a vector back into the shape of `μ`. The layout is a value and +# not a closure, so a layer can store it in a field. `NullParameters` flattens to an empty vector. +_flatten_system_parameters(μ::NamedTuple) = flatten(μ) +_flatten_system_parameters(::NullParameters) = flatten(NamedTuple()) +_flatten_system_parameters(::Type{T}, μ::NamedTuple) where {T} = flatten(T, μ) +_flatten_system_parameters(::Type{T}, ::NullParameters) where {T} = flatten(T, NamedTuple()) + +# The input of a parametric layer: the state with the flattened system parameters below it, in the +# element type of the state. One `μ` goes below every column of a matrix; a vector of them goes one +# below each column, so the columns are concatenated side by side. Without system parameters the state +# is the input as it is. +concatenate_array_with_parameters(x::AbstractArray, ::NullParameters) = x + +function concatenate_array_with_parameters(x::AbstractVector, μ::NamedTuple) + vcat(x, first(_flatten_system_parameters(eltype(x), μ))) +end + +function concatenate_array_with_parameters(x::AbstractMatrix, μ::NamedTuple) + flat = first(_flatten_system_parameters(eltype(x), μ)) + vcat(x, repeat(flat, 1, size(x, 2))) +end + +function concatenate_array_with_parameters(x::AbstractMatrix, μs::AbstractVector) + size(x, 2) == length(μs) || throw(DimensionMismatch( + "the matrix has $(size(x, 2)) columns and there are $(length(μs)) system parameters")) + reduce(hcat, + [concatenate_array_with_parameters(x[:, i], μ) for (i, μ) in enumerate(μs)]) +end + +@doc raw""" + apply_parametric(model, x, μ, ps) + +Apply the `Chain` `model` with the network parameters `ps` to the input `x` and the parameters `μ` +of the system. + +`x` is one sample — a vector, or a `(q, p)` named tuple of vectors — with one `μ`, an +`OptionalParameters`. Or `x` is a batch of `k` samples — a `2d × 1 × k` tensor, or a `(q, p)` named +tuple of `d × 1 × k` tensors — with a vector of `k` values of `μ`, one per sample, as a +[`ParametricDataLoader`](@ref) draws them. The output has the shape of the input. + +The `Chain` receives the pair `(x, μ)`, so its first layer takes one sample and its `μ`, and a layer +that returns such a pair passes `μ` on to the next. A batch is applied sample by sample. + +# Examples + +```jldoctest +using GeometricMachineLearning +using GeometricMachineLearning: apply_parametric +import AbstractNeuralNetworks: AbstractExplicitLayer + +# a layer that adds the system parameter `a` to its input +struct AddParameter <: AbstractExplicitLayer{2, 2} end +(::AddParameter)((x, μ)::Tuple, ps) = x .+ μ.a + +model = Chain(AddParameter()) +ps = NetworkParameters((L1 = NamedTuple(),)) +x = reshape([1.0, 2.0, 3.0, 4.0], 2, 1, 2) +apply_parametric(model, x, [(a = 10.0,), (a = 20.0,)], ps) + +# output + +2×1×2 Array{Float64, 3}: +[:, :, 1] = + 11.0 + 12.0 + +[:, :, 2] = + 23.0 + 24.0 +``` +""" +function apply_parametric(model::Chain, x::AbstractVector, μ::OptionalParameters, + ps::NetworkParameters) + model((x, μ), ps) +end + +function apply_parametric(model::Chain, x::NamedTupleOfArrays{<:Any, (:q, :p)}, + μ::OptionalParameters, ps::NetworkParameters) + model((x, μ), ps) +end + +function apply_parametric(model::Chain, x::AbstractArray{<:Any, 3}, μs::AbstractVector, + ps::NetworkParameters) + check_parametric_batch(size(x), μs) + outputs = [apply_parametric(model, x[:, 1, i], μ, ps) for (i, μ) in enumerate(μs)] + reshape(reduce(hcat, outputs), :, 1, length(μs)) +end + +function apply_parametric(model::Chain, x::NamedTupleOfArrays{<:Any, (:q, :p)}, + μs::AbstractVector, ps::NetworkParameters) + check_parametric_batch(size(x.q), μs) + outputs = [apply_parametric(model, (q = x.q[:, 1, i], p = x.p[:, 1, i]), μ, ps) + for (i, μ) in enumerate(μs)] + (q = reshape(reduce(hcat, [output.q for output in outputs]), :, 1, length(μs)), + p = reshape(reduce(hcat, [output.p for output in outputs]), :, 1, length(μs))) +end + +# A batch holds one time step per sample and one `μ` per sample. +function check_parametric_batch(input_size::NTuple{3, Int}, μs::AbstractVector) + input_size[2] == 1 && input_size[3] == length(μs) || throw(DimensionMismatch( + "a batch for `apply_parametric` has one time step and one system parameter per " * + "sample; the input has $(input_size[2]) time steps and $(input_size[3]) samples, " * + "and there are $(length(μs)) system parameters")) + nothing +end diff --git a/src/pullbacks/zygote_pullback.jl b/src/pullbacks/zygote_pullback.jl index 6848f7629..4377e0e5f 100644 --- a/src/pullbacks/zygote_pullback.jl +++ b/src/pullbacks/zygote_pullback.jl @@ -25,9 +25,10 @@ struct ZygotePullback{NNLT} <: AbstractPullback{NNLT} loss::NNLT end -# One method for a network that is trained on inputs only, and one for a network that is trained on -# an `(input, output)` pair: the two shapes a `NetworkLoss` takes. A method per shape and not a -# union over them, which says nothing about which one arrived. +# One method for a network that is trained on inputs only, one for a network that is trained on +# an `(input, output)` pair, and two for a network that also takes the parameters of the system, +# one `μ` or one per sample: the shapes a `NetworkLoss` takes. A method per shape and not a union +# over them, which says nothing about which one arrived. function (_pullback::ZygotePullback)(ps::NetworkParameters, model, input_nt::ArrayOrNamedTuple) value, pb = Zygote.pullback(ps -> _pullback.loss(model, ps, input_nt), ps) value, Δ -> only(pb(Δ)) @@ -40,3 +41,18 @@ function (_pullback::ZygotePullback)( ps -> _pullback.loss(model, ps, input_nt_output_nt...), ps) value, Δ -> only(pb(Δ)) end + +function (_pullback::ZygotePullback)( + ps::NetworkParameters, model, + input_output_μ::Tuple{ + <:ArrayOrNamedTuple, <:ArrayOrNamedTuple, <:OptionalParameters}) + value, pb = Zygote.pullback(ps -> _pullback.loss(model, ps, input_output_μ...), ps) + value, Δ -> only(pb(Δ)) +end + +function (_pullback::ZygotePullback)( + ps::NetworkParameters, model, + input_output_μs::Tuple{<:ArrayOrNamedTuple, <:ArrayOrNamedTuple, <:AbstractVector}) + value, pb = Zygote.pullback(ps -> _pullback.loss(model, ps, input_output_μs...), ps) + value, Δ -> only(pb(Δ)) +end diff --git a/test/data_loader/parametric_data_loader.jl b/test/data_loader/parametric_data_loader.jl new file mode 100644 index 000000000..0fdfeb1b4 --- /dev/null +++ b/test/data_loader/parametric_data_loader.jl @@ -0,0 +1,158 @@ +using GeometricMachineLearning +using GeometricMachineLearning: ParametricDataLoader, ParametricLoss, params, + convert_input_and_batch_indices_to_array +using GeometricProblems.CoupledHarmonicOscillator: hodeensemble, default_parameters +using GeometricIntegrators: ImplicitMidpoint, integrate +using Test +import Random + +include("../helpers/eltypes.jl") +include("../helpers/parametric_affine.jl") + +# `m` trajectories of `nt` time steps of a state of dimension `2d`. +const d, nt, m = 1, 5, 3 + +# The data of trajectory `j` at time step `t`, in every entry, and its system parameters `μⱼ`. +constant_data(::Type{T}) where {T} = [T(j) for _ in 1:(2d), t in 1:nt, j in 1:m] +timed_data(::Type{T}) where {T} = [T(100j + t) for _ in 1:(2d), t in 1:nt, j in 1:m] +system_parameters(::Type{T}) where {T} = [(a = T(j),) for j in 1:m] + +# The loader of a `2d × nt × m` array, and of the `(q, p)` pair of its two halves. +function loaders(data, μs) + (ParametricDataLoader(data, μs), + ParametricDataLoader((q = data[1:d, :, :], p = data[(d + 1):end, :, :]), μs)) +end + +# Entry `k` of a drawn batch, as the vector of the state it holds. +sample(x::AbstractArray, k) = x[:, 1, k] +sample(x::NamedTuple, k) = vcat(x.q[:, 1, k], x.p[:, 1, k]) + +# Every sample of one epoch: its start position and trajectory, its input, its output and its `μ`. +function epoch(dl, batch) + samples = [] + for indices in batch(dl) + input, output, μs = convert_input_and_batch_indices_to_array(dl, batch, indices) + @test length(μs) == length(indices) + for (k, index) in enumerate(indices) + push!(samples, (index, sample(input, k), sample(output, k), μs[k])) + end + end + samples +end + +all_pairs = sort(vec([(t, j) for t in 1:(nt - 1), j in 1:m])) + +@testset "one epoch of Batch(b) on a ParametricDataLoader ($T)" for T in REAL_ELTYPES + Random.seed!(1) + for dl in loaders(constant_data(T), system_parameters(T)) + samples = epoch(dl, Batch(4)) + # each pair of a start position and a trajectory once + @test sort([index for (index, _, _, _) in samples]) == all_pairs + for (_, input, output, μ) in samples + @test μ.a == input[1] + @test all(==(input[1]), input) + end + # the last minibatch is shorter: 12 pairs in minibatches of 5 + @test length.(Batch(5)(dl)) == [5, 5, 2] + end + for dl in loaders(timed_data(T), system_parameters(T)) + samples = epoch(dl, Batch(4)) + @test sort([index for (index, _, _, _) in samples]) == all_pairs + for ((t, j), input, output, μ) in samples + @test input == fill(T(100j + t), 2d) + # the output is the input one time step later + @test output == input .+ 1 + @test μ.a == T(j) + @test eltype(input) === eltype(output) === T + end + end +end + +@testset "a ParametricDataLoader of an EnsembleSolution" begin + parameters = [default_parameters(), map(v -> v .+ 1, default_parameters())] + solution = integrate(hodeensemble(; parameters = parameters), ImplicitMidpoint()) + dl = ParametricDataLoader(solution) + @test dl.parameters == parameters + @test dl.n_params == 2 + n_steps = length(solution.t) + @test dl.input_time_steps == n_steps + for j in 1:2, t in 1:n_steps + + @test dl.input.q[:, t, j] == solution.s[j].q[t - 1] + @test dl.input.p[:, t, j] == solution.s[j].p[t - 1] + end + Random.seed!(2) + samples = epoch(dl, Batch(7)) + @test sort([index for (index, _, _, _) in samples]) == + sort(vec([(t, j) for t in 1:(n_steps - 1), j in 1:2])) + for ((t, j), input, output, μ) in samples + @test μ == parameters[j] + @test input == vcat(dl.input.q[:, t, j], dl.input.p[:, t, j]) + @test output == vcat(dl.input.q[:, t + 1, j], dl.input.p[:, t + 1, j]) + end +end + +# The message of the `ArgumentError` that `f` throws. +function argument_error_message(f) + try + f() + catch error + error isa ArgumentError && return error.msg + rethrow() + end + "no error" +end + +@testset "the edges of a ParametricDataLoader" begin + data = timed_data(Float64) + μs = system_parameters(Float64) + dl = ParametricDataLoader(data, μs) + + # one time step per sample: a sequence length other than one is an error before any step + message = argument_error_message(() -> Batch(2, 3, 1)(dl)) + @test occursin("3", message) + @test occursin("sequence length", message) + + # a parameter vector of the wrong length, and no trajectory + message = argument_error_message(() -> ParametricDataLoader(data, μs[1:2])) + @test occursin("2", message) && occursin("3", message) + message = argument_error_message( + () -> ParametricDataLoader(zeros(2d, nt, 0), NamedTuple{(:a,), Tuple{Float64}}[])) + @test occursin("0", message) + + # too few time steps for the batch + message = argument_error_message(() -> Batch(2)(ParametricDataLoader(data[:, 1:1, :], μs))) + @test occursin("1 time step", message) + + # one element type for the data and every `μ` + message = argument_error_message(() -> ParametricDataLoader(Float32.(data), μs)) + @test occursin("Float32", message) && occursin("Float64", message) + @test ParametricDataLoader(data, [(a = 1.0, b = [2.0, 3.0]) for _ in 1:m]).n_params == m + # an empty `μ` has no element type, and concatenates nothing + @test ParametricDataLoader(data, [(;) for _ in 1:m]).n_params == m + + # `q` and `p` of one size + message = argument_error_message( + () -> ParametricDataLoader((q = data[1:1, :, :], p = data[:, :, :]), μs)) + @test occursin("(1, $(nt), $(m))", message) && occursin("(2, $(nt), $(m))", message) +end + +# The network is `(x, μ) ↦ W [x; flatten(μ)] + b`. `train!` without `loss` and with +# `loss = ParametricLoss()` draw the same network and the same minibatches under one seed. +@testset "train! on a ParametricDataLoader trains on ParametricLoss ($T)" for T in REAL_ELTYPES + dl = ParametricDataLoader(timed_data(T) ./ 100, system_parameters(T)) + function trained(; kwargs...) + Random.seed!(3) + nn = NeuralNetwork(Chain(ParametricAffine{2d}(1)), T) + initial = deepcopy(params(nn)) + opt = TrainingOptimizer(nn.params; algorithm = Adam()) + history = train!(nn, opt, dl, Batch(4), 3; show_progress = false, kwargs...) + history, initial, params(nn) + end + history, initial, final = trained() + @test history == first(trained(; loss = ParametricLoss())) + @test length(history) == 3 + @test eltype(history) === T + @test final.L1.W != initial.L1.W + @test final.L1.b != initial.L1.b +end diff --git a/test/helpers/parametric_affine.jl b/test/helpers/parametric_affine.jl new file mode 100644 index 000000000..2fac1b5fb --- /dev/null +++ b/test/helpers/parametric_affine.jl @@ -0,0 +1,35 @@ +# A layer whose forward pass takes the parameters `μ` of the system beside the state: +# `(x, μ) ↦ W [x; flatten(μ)] + b`. It acts on one sample, a vector `x` or a `(q, p)` pair of vectors +# that stacks `q` above `p`, which is what `apply_parametric` hands the `Chain`. The tests build the +# expected value of this map by hand, with the flattened `μ` written out. +using GeometricMachineLearning: concatenate_array_with_parameters +using AbstractNeuralNetworks: AbstractExplicitLayer, Initializer, NamedTupleOfArrays +import AbstractNeuralNetworks: initialparameters, parameterlength +import KernelAbstractions +import Random + +struct ParametricAffine{M} <: AbstractExplicitLayer{M, M} + parameter_length::Int +end + +function initialparameters( + rng::Random.AbstractRNG, init::Initializer, l::ParametricAffine{M}, + backend::KernelAbstractions.Backend, ::Type{T}) where {M, T} + W = KernelAbstractions.zeros(backend, T, M, M + l.parameter_length) + b = KernelAbstractions.zeros(backend, T, M) + init(rng, W) + init(rng, b) + (W = W, b = b) +end + +parameterlength(l::ParametricAffine{M}) where {M} = M * (M + l.parameter_length) + M + +(l::ParametricAffine)(input::Tuple, ps::NamedTuple) = affine(input..., ps) + +affine(x::AbstractVector, μ, ps) = ps.W * concatenate_array_with_parameters(x, μ) + ps.b + +function affine(x::NamedTupleOfArrays{<:Any, (:q, :p)}, μ, ps) + y = affine(vcat(x.q, x.p), μ, ps) + n = length(x.q) + (q = y[1:n], p = y[(n + 1):end]) +end diff --git a/test/loss/parametric_loss.jl b/test/loss/parametric_loss.jl new file mode 100644 index 000000000..07eee3f16 --- /dev/null +++ b/test/loss/parametric_loss.jl @@ -0,0 +1,82 @@ +using GeometricMachineLearning +using GeometricMachineLearning: ParametricLoss, ZygotePullback +using FiniteDifferences: central_fdm, grad +using LinearAlgebra: norm +using Test +import Random + +include("../helpers/eltypes.jl") +include("../helpers/parametric_affine.jl") + +# System parameters in `Float64`, with a vector-valued leaf, and their flattened form written out by +# hand in the element type `T`. +μ(j) = (a = Float64(j), b = [0.5j, -1.5j]) +flat(::Type{T}, μ) where {T} = T[μ.a; μ.b[1]; μ.b[2]] + +const K = 4 +const μS = [μ(j) for j in 1:K] + +# The network `(x, μ) ↦ W [x; flatten(μ)] + b` on a batch of `K` samples, and the relative norm of its +# error, both written out with the flattened `μ` built by hand. +function closed_form(W, b, X, μs) + T = eltype(X) + reshape(reduce(hcat, [W * vcat(X[:, 1, i], flat(T, μs[i])) + b for i in axes(X, 3)]), + size(X)) +end +closed_form_loss(W, b, X, Y, μs) = norm(closed_form(W, b, X, μs) - Y) / norm(Y) + +function data(::Type{T}) where {T} + rng = Random.Xoshiro(7) + W = randn(rng, T, 2, 5) + b = randn(rng, T, 2) + X = randn(rng, T, 2, 1, K) + Y = randn(rng, T, 2, 1, K) + W, b, X, Y +end + +network_parameters(W, b) = NetworkParameters((L1 = (W = W, b = b),)) +const MODEL = Chain(ParametricAffine{2}(3)) + +@testset "ParametricLoss is the relative norm of the error ($T)" for T in REAL_ELTYPES + W, b, X, Y = data(T) + value = ParametricLoss()(MODEL, network_parameters(W, b), X, Y, μS) + @test value isa T + @test isapprox(value, closed_form_loss(W, b, X, Y, μS); rtol = 10 * eps(T)) +end + +@testset "ZygotePullback of ParametricLoss against finite differences" begin + W, b, X, Y = data(Float64) + value, back = ZygotePullback(ParametricLoss())( + network_parameters(W, b), MODEL, (X, Y, μS)) + @test isapprox(value, closed_form_loss(W, b, X, Y, μS); rtol = 10 * eps(Float64)) + + gradient = back(one(value)) + # the finite-difference gradient of the closed form over the entries of `W`, then of `b` + loss_of_entries(v) = closed_form_loss(reshape(v[1:10], 2, 5), v[11:12], X, Y, μS) + reference = only(grad(central_fdm(5, 1), loss_of_entries, vcat(vec(W), b))) + computed = vcat(vec(gradient.L1.W), gradient.L1.b) + @test norm(computed - reference) ≤ 1e-6 * norm(reference) +end + +# One sample with one `μ` is the other shape the pullback takes. +@testset "ZygotePullback of ParametricLoss on one sample ($T)" for T in REAL_ELTYPES + W, b, X, Y = data(T) + x, y = X[:, 1, 1], Y[:, 1, 1] + value, back = ZygotePullback(ParametricLoss())( + network_parameters(W, b), MODEL, (x, y, μ(2))) + expected = norm(W * vcat(x, flat(T, μ(2))) + b - y) / norm(y) + @test value isa T + @test isapprox(value, expected; rtol = 10 * eps(T)) + @test eltype(back(one(value)).L1.W) === T +end + +@testset "ZygotePullback of ParametricLoss keeps Float32" begin + W, b, X, Y = data(Float32) + value, back = ZygotePullback(ParametricLoss())( + network_parameters(W, b), MODEL, (X, Y, μS)) + @test value isa Float32 + @test isapprox(value, closed_form_loss(W, b, X, Y, μS); rtol = 10 * eps(Float32)) + gradient = back(one(value)) + @test eltype(gradient.L1.W) === Float32 + @test eltype(gradient.L1.b) === Float32 +end diff --git a/test/parametric.jl b/test/parametric.jl new file mode 100644 index 000000000..66dfcd5eb --- /dev/null +++ b/test/parametric.jl @@ -0,0 +1,94 @@ +using GeometricMachineLearning +using GeometricMachineLearning: concatenate_array_with_parameters, apply_parametric, + NullParameters +using Test +import Random + +include("helpers/eltypes.jl") +include("helpers/parametric_affine.jl") + +# System parameters in `Float64`, with a vector-valued leaf; `flat` writes out by hand what +# flattening one of them gives in the element type `T`: the leaves in the order of the keys. +μ(j) = (a = Float64(j), b = [0.5j, -1.5j]) +flat(::Type{T}, μ) where {T} = T[μ.a; μ.b[1]; μ.b[2]] + +@testset "concatenate_array_with_parameters ($T)" for T in REAL_ELTYPES + x = T[1 2 3; 4 5 6] + + # one `μ` per column: the columns side by side + concatenated = concatenate_array_with_parameters(x, [μ(1), μ(2), μ(3)]) + @test concatenated == T[1 2 3; 4 5 6; 1 2 3; 0.5 1 1.5; -1.5 -3 -4.5] + @test eltype(concatenated) === T + + # one `μ` below a vector + concatenated = concatenate_array_with_parameters(x[:, 2], μ(2)) + @test concatenated == T[2, 5, 2, 1, -3] + @test eltype(concatenated) === T + + # one `μ` below every column of a matrix + concatenated = concatenate_array_with_parameters(x, μ(2)) + @test concatenated == T[1 2 3; 4 5 6; 2 2 2; 1 1 1; -3 -3 -3] + @test eltype(concatenated) === T + + # no system parameters: the state itself + @test concatenate_array_with_parameters(x, NullParameters()) === x + @test concatenate_array_with_parameters(x[:, 1], NullParameters()) == x[:, 1] + # an empty named tuple concatenates nothing + @test concatenate_array_with_parameters(x[:, 1], (;)) == x[:, 1] + @test eltype(concatenate_array_with_parameters(x[:, 1], (;))) === T + + @test_throws DimensionMismatch concatenate_array_with_parameters(x, [μ(1), μ(2)]) +end + +# `ParametricAffine` maps `(x, μ)` to `W [x; flatten(μ)] + b`; `closed_form` writes that out with the +# flattened `μ` built by hand. +@testset "apply_parametric against the closed form ($T)" for T in REAL_ELTYPES + W = T[0.5 -1 0.25 2 -0.5; 1.5 0.75 -2 0.125 1] + b = T[0.1, -0.3] + ps = NetworkParameters((L1 = (W = W, b = b),)) + model = Chain(ParametricAffine{2}(3)) + closed_form(x, μ) = W * vcat(x, flat(T, μ)) + b + tolerance = 10 * eps(T) + + # one sample: a vector, and a `(q, p)` pair of vectors + x = T[0.3, -0.7] + output = apply_parametric(model, x, μ(2), ps) + @test size(output) == size(x) + @test eltype(output) === T + @test isapprox(output, closed_form(x, μ(2)); rtol = tolerance) + + output = apply_parametric(model, (q = x[1:1], p = x[2:2]), μ(2), ps) + @test size(output.q) == size(output.p) == (1,) + @test eltype(output.q) === eltype(output.p) === T + @test isapprox(vcat(output.q, output.p), closed_form(x, μ(2)); rtol = tolerance) + + # a batch of `k` samples with `k` values of `μ`: a `2d × 1 × k` tensor, and a `(q, p)` pair of + # `d × 1 × k` tensors + k = 4 + μs = [μ(j) for j in 1:k] + X = reshape(T[0.3, -0.7, 1.1, 0.2, -0.4, 0.9, 0.0, 2.5], 2, 1, k) + expected = reshape(reduce(hcat, [closed_form(X[:, 1, i], μs[i]) for i in 1:k]), 2, 1, k) + + output = apply_parametric(model, X, μs, ps) + @test size(output) == size(X) + @test eltype(output) === T + @test isapprox(output, expected; rtol = tolerance) + + output = apply_parametric(model, (q = X[1:1, :, :], p = X[2:2, :, :]), μs, ps) + @test size(output.q) == size(output.p) == (1, 1, k) + @test eltype(output.q) === eltype(output.p) === T + @test isapprox(output.q, expected[1:1, :, :]; rtol = tolerance) + @test isapprox(output.p, expected[2:2, :, :]; rtol = tolerance) + + # one `μ` per sample and one time step per sample + @test_throws DimensionMismatch apply_parametric(model, X, μs[1:3], ps) + @test_throws DimensionMismatch apply_parametric( + model, reshape(X, 2, 2, 2), μs[1:2], ps) +end + +@testset "the parametric names are public and not exported" begin + for name in (:apply_parametric, :ParametricLoss, :ParametricDataLoader) + @test Base.ispublic(GeometricMachineLearning, name) + @test !Base.isexported(GeometricMachineLearning, name) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 8a25eef3d..a7e0b11ac 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -70,6 +70,8 @@ if "core" in GROUPS @safetestset "Symplectic attention: NetworkParameters gradient structure" include("layers/symplectic_attention_network_parameters_gradient.jl") @safetestset "ZygotePullback returns a NetworkParameters" include("pullbacks/zygote_pullback.jl") @safetestset "AutoEncoderLoss and a loss's NetworkParameters" include("loss/losses.jl") + @safetestset "apply_parametric and the concatenation of the system parameters" include("parametric.jl") + @safetestset "ParametricLoss and its ZygotePullback" include("loss/parametric_loss.jl") @safetestset "map_to_cpu" include("map_to_cpu.jl") @safetestset "changebackend" include("integration/changebackend_tests.jl") @safetestset "HDF5 save/load for GML special array types" include("integration/hdf5_support.jl") @@ -87,6 +89,7 @@ if "core" in GROUPS @safetestset "Classifier accuracy" include("data_loader/accuracy.jl") @safetestset "DataLoader batching" include("data_loader/batch.jl") @safetestset "DataLoader constructors" include("data_loader/data_loader.jl") + @safetestset "ParametricDataLoader batching and training" include("data_loader/parametric_data_loader.jl") @safetestset "Data loader docstring examples" include("integration/docstrings/data_loader.jl") @safetestset "Layer and architecture docstring examples" include("integration/docstrings/layers_and_architectures.jl") @safetestset "Loss docstring examples" include("integration/docstrings/losses.jl") From 7709f785778b084797a980e4de93958465e8d4ff Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Sun, 25 Oct 2026 19:49:01 +0100 Subject: [PATCH 2/7] =?UTF-8?q?Refuse=20a=20mixed=20(q,=20p)=20pair,=20a?= =?UTF-8?q?=20prediction=20window=20other=20than=20one=20and=20mixed=20?= =?UTF-8?q?=CE=BC=20shapes=20in=20ParametricDataLoader?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The (q, p) constructor takes one array type, as gather does; a pair of two types is a MethodError at construction instead of in the first minibatch. The batch functor throws an ArgumentError for a prediction window other than one, and the constructor one for a μ whose parameterlayout differs from the first. Drop the unread input_dim field and merge the two single-sample apply_parametric methods. Record K21 (empty batch) and K22 (upstream Float32 promotion), and extend K20. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 +- KNOWN_ISSUES.md | 37 ++++++++++++++++++++++ src/GeometricMachineLearning.jl | 5 +-- src/data_loader/parametric_data_loader.jl | 35 ++++++++++++-------- src/parametric.jl | 8 ++--- test/data_loader/parametric_data_loader.jl | 21 ++++++++++++ 6 files changed, 85 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2956eef92..686eefa60 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -99,7 +99,7 @@ breaking release). ### Added -- **`apply_parametric`, `ParametricLoss` and `ParametricDataLoader` train a network whose forward pass takes the parameters of the system beside its input.** The three names are `public` and not exported, so a caller qualifies them or names them in a `using GeometricMachineLearning: …`. `apply_parametric(model, x, μ, ps)` applies a `Chain` to one sample — a vector or a `(q, p)` pair of vectors — with one `μ`, or to a batch — a `2d × 1 × k` tensor or a `(q, p)` pair of `d × 1 × k` tensors — with a vector of `k` values of `μ`, sample by sample; the output has the shape of the input. The `Chain` receives the pair `(x, μ)`, and the internal `concatenate_array_with_parameters` puts the flattened `μ` below the state in the element type of the state, side by side for a matrix and a vector of `μ`, and repeated in every column for a matrix and one `μ`. `ParametricLoss()(model, ps, input, output, μ)` is the relative norm of the error of `apply_parametric`, and `ZygotePullback` takes the triple `(input, output, μ)` with one `μ` or a vector of them. `ParametricDataLoader(data, parameters)` takes a `2d × nt × m` array or a `(q, p)` pair of `d × nt × m` arrays with a vector of `m` named tuples, and `ParametricDataLoader(ensemble_solution)` takes an `EnsembleSolution` of `(q, p)` states with the parameters of its problem. A `Batch(b)` draws each pair of a start position and a trajectory once per epoch, as for `DataLoader`, and a minibatch is the state, the state one time step later and the `μ` of each sample. `train!` takes a `ParametricDataLoader` and trains on `ParametricLoss()` when no `loss` is given. A `Batch` with a sequence length other than one, a parameter vector whose length is not the number of trajectories, no trajectory, too few time steps for the batch, and a `μ` whose element type is not the data's raise an `ArgumentError`. The code is ported from pull request #207 by Benedikt Brantner (relates to #207); there `concatenate_array_with_parameters` promoted a `Float32` state to `Float64`, rejected a matrix of more than one column with one `μ`, and the loader and the batch checked their counts with `@assert`. +- **`apply_parametric`, `ParametricLoss` and `ParametricDataLoader` train a network whose forward pass takes the parameters of the system beside its input.** The three names are `public` and not exported, so a caller qualifies them or names them in a `using GeometricMachineLearning: …`. `apply_parametric(model, x, μ, ps)` applies a `Chain` to one sample — a vector or a `(q, p)` pair of vectors — with one `μ`, or to a batch — a `2d × 1 × k` tensor or a `(q, p)` pair of `d × 1 × k` tensors — with a vector of `k` values of `μ`, sample by sample; the output has the shape of the input. The `Chain` receives the pair `(x, μ)`, and the internal `concatenate_array_with_parameters` puts the flattened `μ` below the state in the element type of the state, side by side for a matrix and a vector of `μ`, and repeated in every column for a matrix and one `μ`. `ParametricLoss()(model, ps, input, output, μ)` is the relative norm of the error of `apply_parametric`, and `ZygotePullback` takes the triple `(input, output, μ)` with one `μ` or a vector of them. `ParametricDataLoader(data, parameters)` takes a `2d × nt × m` array or a `(q, p)` pair of `d × nt × m` arrays of one array type with a vector of `m` named tuples, and `ParametricDataLoader(ensemble_solution)` takes an `EnsembleSolution` of `(q, p)` states with the parameters of its problem. A `Batch(b)` draws each pair of a start position and a trajectory once per epoch, as for `DataLoader`, and a minibatch is the state, the state one time step later and the `μ` of each sample. `train!` takes a `ParametricDataLoader` and trains on `ParametricLoss()` when no `loss` is given. A `Batch` with a sequence length or a prediction window other than one, a parameter vector whose length is not the number of trajectories, no trajectory, too few time steps for the batch, a `μ` whose element type is not the data's, and a `μ` whose keys, their order or the sizes of its leaves differ from the first `μ`'s raise an `ArgumentError`; a `(q, p)` pair of two array types is a `MethodError`, as for `DataLoader` (K20). The code is ported from pull request #207 by Benedikt Brantner (relates to #207); there `concatenate_array_with_parameters` promoted a `Float32` state to `Float64`, rejected a matrix of more than one column with one `μ`, and the loader and the batch checked their counts with `@assert`. - **`KernelAbstractions.get_backend` is defined for `PoissonTensor`.** GeometricOptimizers 0.9 asks every operand of a checked product for its backend, so `𝕁 * Y` with a `StiefelManifold` raised `ArgumentError: Implement KernelAbstractions.get_backend(::PoissonTensor...)` without it. diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index f476dae34..bf4e7a13a 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -129,6 +129,26 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. The same call with `q = A` gives a `DataLoader` with `input_dim == 4`. The docstring at `src/data_loader/data_loader.jl:205` states the one-type signature. Closing it means widening the three signatures to `NamedTupleOfArrays{T, (:q, :p)}`, and the docstring with them. + `ParametricDataLoader`'s `(q, p)` constructor has the same one-type restriction and widens with + it. + +### K21 · `concatenate_array_with_parameters` and `apply_parametric` raise Base's empty-reduction error on a batch of zero samples. + +- **location:** `src/parametric.jl` +- **kind:** found late +- **found:** 2026-10-25, by the critics of the parametric core +- **evidence:** + Both reduce the samples with `reduce(hcat, …)` and no `init`. No `ParametricDataLoader` draws an + empty minibatch, and no `init` gives a width that holds for every network. A direct call fails: + + ```julia + using GeometricMachineLearning: concatenate_array_with_parameters + concatenate_array_with_parameters(zeros(Float32, 2, 0), NamedTuple{(:a,), Tuple{Float32}}[]) + # ArgumentError: reducing over an empty collection is not allowed; … + ``` + + `apply_parametric` on a `2 × 1 × 0` tensor with no `μ` raises the same error. A `2 × 0` matrix + with one `μ` gives a `3 × 0` matrix. ## Upstream @@ -145,6 +165,23 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. The blocks are not test failures: both runs are green, and every testset of ea0c397 has the same count in both. +### K22 · AbstractNeuralNetworks' loss norm of a `(q, p)` named tuple promotes `Float32` to `Float64`. + +- **location:** AbstractNeuralNetworks 0.9.3, `src/losses.jl:73` +- **kind:** upstream +- **found:** 2026-10-25, by the critics of the parametric core +- **evidence:** + `_norm(dx::NamedTuple) = sum(map(_norm, values(dx))) / √length(dx)` divides by the square root of + an `Int`, which is a `Float64`. So `ParametricLoss` returns a `Float64` on `Float32` `(q, p)` + data, and `train!` on a `Float32` `(q, p)` `ParametricDataLoader` returns a `Vector{Float64}` + history. Array data keep `Float32`. + + ```julia + import AbstractNeuralNetworks + x = (q = rand(Float32, 2, 1, 3), p = rand(Float32, 2, 1, 3)) + typeof(AbstractNeuralNetworks._norm(x)) # Float64 + ``` + ### K4 · Metal.jl 1.11.1 needs a Metal device to precompile, so a sandboxed run with no Metal cache fails before the tests - **location:** `test/Project.toml:16` diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index 6390dd02c..9d61e4901 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -12,8 +12,9 @@ using AbstractNeuralNetworks import NeuralNetworkParameters: NetworkParameters using NeuralNetworkParameters: mapstorage # `flatten` puts the parameters of a system into the input of a network that takes them -# (`src/parametric.jl`), and `parameter_eltype` checks their element type against the data's. -using NeuralNetworkParameters: flatten, parameter_eltype +# (`src/parametric.jl`); `parameter_eltype` checks their element type against the data's, and +# `parameterlayout` their shape against each other's. +using NeuralNetworkParameters: flatten, parameter_eltype, parameterlayout using ChainRulesCore using GeometricBase using GeometricSolutions: GeometricSolution, EnsembleSolution, TimeSeries diff --git a/src/data_loader/parametric_data_loader.jl b/src/data_loader/parametric_data_loader.jl index ea51f6eb8..640a24b6e 100644 --- a/src/data_loader/parametric_data_loader.jl +++ b/src/data_loader/parametric_data_loader.jl @@ -5,11 +5,12 @@ Make a loader of trajectories of a system together with the parameters ``\mu`` of the system that each trajectory belongs to. -`data` is a `2d × nt × m` array, or a `(q, p)` named tuple of two `d × nt × m` arrays: `m` -trajectories of `nt` time steps, with `q` above `p` where they are stacked. `parameters` is a vector -of `m` named tuples, the ``\mu`` of each trajectory. The data and every ``\mu`` have one element type. -An `EnsembleSolution` of a problem with `(q, p)` states gives its `q` and `p` and the parameters of its -problem. +`data` is a `2d × nt × m` array, or a `(q, p)` named tuple of two `d × nt × m` arrays of one +array type: `m` trajectories of `nt` time steps, with `q` above `p` where they are stacked. +`parameters` is a vector of `m` named tuples, the ``\mu`` of each trajectory. Every ``\mu`` has the +keys of the first, in the same order and at every level, and the same sizes of its leaves. The data +and every ``\mu`` have one element type. An `EnsembleSolution` of a problem with `(q, p)` states +gives its `q` and `p` and the parameters of its problem. # The batch @@ -18,8 +19,8 @@ minibatches of `batch_size`; the last minibatch is shorter. A drawn minibatch is `(input, output, μs)`: the state at the start position, the state one time step later, and the ``\mu`` of each sample. [`train!`](@ref) on a `ParametricDataLoader` trains on [`ParametricLoss`](@ref)`()` when no `loss` is given, so the network is applied with -[`apply_parametric`](@ref), one time step per sample. A `Batch` with a sequence length other than one -is an `ArgumentError`. +[`apply_parametric`](@ref), one time step per sample. A `Batch` with a sequence length or a +prediction window other than one is an `ArgumentError`. A problem without system parameters is trained with a [`DataLoader`](@ref). """ @@ -27,28 +28,31 @@ struct ParametricDataLoader{ T, AT <: ArrayOrNamedTuple{T}, PT <: AbstractVector{<:NamedTuple}} input::AT parameters::PT - input_dim::Int input_time_steps::Int n_params::Int function ParametricDataLoader(input::AT, parameters::PT) where {T, AT <: Union{AbstractArray{T, 3}, - NamedTuple{(:q, :p), <:Tuple{AbstractArray{T, 3}, AbstractArray{T, 3}}}}, + NamedTuple{(:q, :p), Tuple{AT3, AT3}} where {AT3 <: AbstractArray{T, 3}}}, PT <: AbstractVector{<:NamedTuple}} - input_dim, input_time_steps, n_params = data_size(input) + _, input_time_steps, n_params = data_size(input) n_params ≥ 1 && length(parameters) == n_params || throw(ArgumentError( "the data hold $(n_params) trajectories and there are $(length(parameters)) " * "system parameters; a `ParametricDataLoader` takes one for each trajectory, " * "and at least one trajectory.")) - for μ in parameters + layout = parameterlayout(first(parameters)) + for (j, μ) in enumerate(parameters) S = parameter_eltype(μ) S === T || S === Union{} || throw(ArgumentError( "the data have the element type $(T) and a system parameter has the " * "element type $(S); a `ParametricDataLoader` takes one element type.")) + parameterlayout(μ) == layout || throw(ArgumentError( + "system parameter $(j) differs from system parameter 1 in its keys, their " * + "order or the sizes of its leaves; a `ParametricDataLoader` takes one shape.")) end - new{T, AT, PT}(input, parameters, input_dim, input_time_steps, n_params) + new{T, AT, PT}(input, parameters, input_time_steps, n_params) end end @@ -57,7 +61,7 @@ data_size(input::AbstractArray{<:Any, 3}) = size(input) function data_size(input::NamedTuple{(:q, :p)}) size(input.q) == size(input.p) || throw(ArgumentError( "`q` has the size $(size(input.q)) and `p` has the size $(size(input.p)).")) - (2 * size(input.q, 1), size(input.q, 2), size(input.q, 3)) + size(input.q) end function ParametricDataLoader(ensemble_solution::EnsembleSolution{T, @@ -79,11 +83,14 @@ function networkbackend(dl::ParametricDataLoader{T, <:NamedTuple}) where {T} end # The start positions are those with a target `prediction_window` steps later; the network takes one -# time step per sample, so the sequence length is one. +# time step per sample and returns one, so the sequence length and the prediction window are one. function (batch::Batch)(dl::ParametricDataLoader) batch.seq_length == 1 || throw(ArgumentError( "a `ParametricDataLoader` draws one time step per sample, and the `Batch` asks for a " * "sequence length of $(batch.seq_length).")) + batch.prediction_window == 1 || throw(ArgumentError( + "a `ParametricDataLoader` draws a target of one time step per sample, and the `Batch` " * + "asks for a prediction window of $(batch.prediction_window).")) n = dl.input_time_steps - batch.prediction_window n ≥ 1 || throw(ArgumentError( "the `ParametricDataLoader` has $(dl.input_time_steps) time steps and the `Batch` asks " * diff --git a/src/parametric.jl b/src/parametric.jl index d11fd7510..f41a129d1 100644 --- a/src/parametric.jl +++ b/src/parametric.jl @@ -74,12 +74,8 @@ apply_parametric(model, x, [(a = 10.0,), (a = 20.0,)], ps) 24.0 ``` """ -function apply_parametric(model::Chain, x::AbstractVector, μ::OptionalParameters, - ps::NetworkParameters) - model((x, μ), ps) -end - -function apply_parametric(model::Chain, x::NamedTupleOfArrays{<:Any, (:q, :p)}, +function apply_parametric(model::Chain, + x::Union{AbstractVector, NamedTupleOfArrays{<:Any, (:q, :p)}}, μ::OptionalParameters, ps::NetworkParameters) model((x, μ), ps) end diff --git a/test/data_loader/parametric_data_loader.jl b/test/data_loader/parametric_data_loader.jl index 0fdfeb1b4..ffbfe74e7 100644 --- a/test/data_loader/parametric_data_loader.jl +++ b/test/data_loader/parametric_data_loader.jl @@ -112,6 +112,9 @@ end message = argument_error_message(() -> Batch(2, 3, 1)(dl)) @test occursin("3", message) @test occursin("sequence length", message) + # a prediction window other than one is an error before any step, with the time steps to spare + message = argument_error_message(() -> Batch{:FeedForward}(2, 1, 2)(dl)) + @test occursin("prediction window of 2", message) # a parameter vector of the wrong length, and no trajectory message = argument_error_message(() -> ParametricDataLoader(data, μs[1:2])) @@ -135,6 +138,24 @@ end message = argument_error_message( () -> ParametricDataLoader((q = data[1:1, :, :], p = data[:, :, :]), μs)) @test occursin("(1, $(nt), $(m))", message) && occursin("(2, $(nt), $(m))", message) + + # `q` and `p` of one array type, as `gather` batches them + @test_throws MethodError ParametricDataLoader( + (q = data[1:d, :, :], p = view(data, (d + 1):(2d), :, :)), μs) + views = ParametricDataLoader( + (q = view(data, 1:d, :, :), p = view(data, (d + 1):(2d), :, :)), μs) + input, output, _ = convert_input_and_batch_indices_to_array(views, Batch(2), [(1, 2)]) + @test sample(input, 1) == data[:, 1, 2] && sample(output, 1) == data[:, 2, 2] + + # every `μ` with the keys of the first, in its order and at every level, and its leaf sizes + for parameters in ([(a = 1.0,), (b = 2.0,), (a = 3.0,)], + [(a = 1.0, b = 2.0), (b = 3.0, a = 4.0), (a = 5.0, b = 6.0)], + [(a = 1.0,), (a = 2.0, b = 3.0), (a = 3.0,)], + [(a = (c = 1.0,),), (a = (e = 2.0,),), (a = (c = 3.0,),)], + [(a = [1.0, 2.0],), (a = [1.0, 2.0, 3.0],), (a = [1.0, 2.0],)]) + message = argument_error_message(() -> ParametricDataLoader(data, parameters)) + @test occursin("system parameter 2", message) + end end # The network is `(x, μ) ↦ W [x; flatten(μ)] + b`. `train!` without `loss` and with From 24be74bd00eb033f9fe568b4d75339e081af079c Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Sun, 25 Oct 2026 21:04:37 +0100 Subject: [PATCH 3/7] Map the batches of apply_parametric over columns, narrow its one-sample method, and drop K21 Co-Authored-By: Claude Opus 5.5 --- KNOWN_ISSUES.md | 18 ------ src/data_loader/parametric_data_loader.jl | 7 ++- src/parametric.jl | 17 ++++-- test/parametric.jl | 68 ++++++++++++++++++++++- 4 files changed, 83 insertions(+), 27 deletions(-) diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index bf4e7a13a..7219b932d 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -132,24 +132,6 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. `ParametricDataLoader`'s `(q, p)` constructor has the same one-type restriction and widens with it. -### K21 · `concatenate_array_with_parameters` and `apply_parametric` raise Base's empty-reduction error on a batch of zero samples. - -- **location:** `src/parametric.jl` -- **kind:** found late -- **found:** 2026-10-25, by the critics of the parametric core -- **evidence:** - Both reduce the samples with `reduce(hcat, …)` and no `init`. No `ParametricDataLoader` draws an - empty minibatch, and no `init` gives a width that holds for every network. A direct call fails: - - ```julia - using GeometricMachineLearning: concatenate_array_with_parameters - concatenate_array_with_parameters(zeros(Float32, 2, 0), NamedTuple{(:a,), Tuple{Float32}}[]) - # ArgumentError: reducing over an empty collection is not allowed; … - ``` - - `apply_parametric` on a `2 × 1 × 0` tensor with no `μ` raises the same error. A `2 × 0` matrix - with one `μ` gives a `3 × 0` matrix. - ## Upstream ### K2 · Revise prints EMFILE errors in the test log diff --git a/src/data_loader/parametric_data_loader.jl b/src/data_loader/parametric_data_loader.jl index 640a24b6e..8353d640e 100644 --- a/src/data_loader/parametric_data_loader.jl +++ b/src/data_loader/parametric_data_loader.jl @@ -36,7 +36,7 @@ struct ParametricDataLoader{ AT <: Union{AbstractArray{T, 3}, NamedTuple{(:q, :p), Tuple{AT3, AT3}} where {AT3 <: AbstractArray{T, 3}}}, PT <: AbstractVector{<:NamedTuple}} - _, input_time_steps, n_params = data_size(input) + input_time_steps, n_params = data_size(input) n_params ≥ 1 && length(parameters) == n_params || throw(ArgumentError( "the data hold $(n_params) trajectories and there are $(length(parameters)) " * "system parameters; a `ParametricDataLoader` takes one for each trajectory, " * @@ -56,12 +56,13 @@ struct ParametricDataLoader{ end end -data_size(input::AbstractArray{<:Any, 3}) = size(input) +# The number of time steps and the number of trajectories of the data. +data_size(input::AbstractArray{<:Any, 3}) = (size(input, 2), size(input, 3)) function data_size(input::NamedTuple{(:q, :p)}) size(input.q) == size(input.p) || throw(ArgumentError( "`q` has the size $(size(input.q)) and `p` has the size $(size(input.p)).")) - size(input.q) + data_size(input.q) end function ParametricDataLoader(ensemble_solution::EnsembleSolution{T, diff --git a/src/parametric.jl b/src/parametric.jl index f41a129d1..5f3a6b02d 100644 --- a/src/parametric.jl +++ b/src/parametric.jl @@ -41,7 +41,8 @@ of the system. `x` is one sample — a vector, or a `(q, p)` named tuple of vectors — with one `μ`, an `OptionalParameters`. Or `x` is a batch of `k` samples — a `2d × 1 × k` tensor, or a `(q, p)` named tuple of `d × 1 × k` tensors — with a vector of `k` values of `μ`, one per sample, as a -[`ParametricDataLoader`](@ref) draws them. The output has the shape of the input. +[`ParametricDataLoader`](@ref) draws them. A batch holds at least one sample. The output has the +shape of the input. The `Chain` receives the pair `(x, μ)`, so its first layer takes one sample and its `μ`, and a layer that returns such a pair passes `μ` on to the next. A batch is applied sample by sample. @@ -75,27 +76,33 @@ apply_parametric(model, x, [(a = 10.0,), (a = 20.0,)], ps) ``` """ function apply_parametric(model::Chain, - x::Union{AbstractVector, NamedTupleOfArrays{<:Any, (:q, :p)}}, + x::Union{ + AbstractVector, NamedTuple{(:q, :p), <:Tuple{AbstractVector, AbstractVector}}}, μ::OptionalParameters, ps::NetworkParameters) model((x, μ), ps) end +# A batch maps over the columns of its `d × k` reshape, not over indexed samples: under Zygote, the +# adjoint of each `getindex` allocates a cotangent of the whole batch, so indexed samples cost O(k²). function apply_parametric(model::Chain, x::AbstractArray{<:Any, 3}, μs::AbstractVector, ps::NetworkParameters) check_parametric_batch(size(x), μs) - outputs = [apply_parametric(model, x[:, 1, i], μ, ps) for (i, μ) in enumerate(μs)] + outputs = map((xᵢ, μ) -> apply_parametric(model, xᵢ, μ, ps), sample_columns(x), μs) reshape(reduce(hcat, outputs), :, 1, length(μs)) end function apply_parametric(model::Chain, x::NamedTupleOfArrays{<:Any, (:q, :p)}, μs::AbstractVector, ps::NetworkParameters) check_parametric_batch(size(x.q), μs) - outputs = [apply_parametric(model, (q = x.q[:, 1, i], p = x.p[:, 1, i]), μ, ps) - for (i, μ) in enumerate(μs)] + outputs = map((q, p, μ) -> apply_parametric(model, (q = q, p = p), μ, ps), + sample_columns(x.q), sample_columns(x.p), μs) (q = reshape(reduce(hcat, [output.q for output in outputs]), :, 1, length(μs)), p = reshape(reduce(hcat, [output.p for output in outputs]), :, 1, length(μs))) end +# The samples of a `d × 1 × k` batch, as the `k` columns of its `d × k` reshape. +sample_columns(x::AbstractArray{<:Any, 3}) = eachcol(reshape(x, size(x, 1), size(x, 3))) + # A batch holds one time step per sample and one `μ` per sample. function check_parametric_batch(input_size::NTuple{3, Int}, μs::AbstractVector) input_size[2] == 1 && input_size[3] == length(μs) || throw(DimensionMismatch( diff --git a/test/parametric.jl b/test/parametric.jl index 66dfcd5eb..aa54b238d 100644 --- a/test/parametric.jl +++ b/test/parametric.jl @@ -1,8 +1,11 @@ using GeometricMachineLearning using GeometricMachineLearning: concatenate_array_with_parameters, apply_parametric, - NullParameters + NullParameters, ParametricLoss, ZygotePullback +using FiniteDifferences: central_fdm, grad +using LinearAlgebra: norm using Test import Random +import Zygote include("helpers/eltypes.jl") include("helpers/parametric_affine.jl") @@ -86,6 +89,69 @@ end model, reshape(X, 2, 2, 2), μs[1:2], ps) end +# The one-sample method takes a vector or a `(q, p)` pair of vectors; a pair of tensors with one `μ` +# matches no method. +@testset "apply_parametric takes one sample as vectors" begin + ps = NetworkParameters((L1 = (W = [0.5 -1 0.25; 1.5 0.75 -2], b = [0.1, -0.3]),)) + model = Chain(ParametricAffine{2}(1)) + X = reshape([0.3, -0.7, 1.1, 0.2], 2, 1, 2) + @test applicable(apply_parametric, model, (q = [0.3], p = [-0.7]), (a = 2.0,), ps) + @test !applicable( + apply_parametric, model, (q = X[1:1, :, :], p = X[2:2, :, :]), (a = 2.0,), ps) +end + +# The bytes that the pullback of `ParametricLoss` and the call of its closure allocate on `data`: the +# least of three readings after two calls that compile. +function pullback_bytes(model, ps, data) + function run() + value, back = ZygotePullback(ParametricLoss())(ps, model, data) + back(one(value)) + end + run() + run() + minimum(@allocated(run()) for _ in 1:3) +end + +# A batch of `k` samples costs O(k): the bytes of one more sample, between 64 and 1024 samples, are +# at most twice the bytes per sample at 64. A sample indexed one at a time in the differentiated path +# costs a cotangent of the whole batch, O(k²) in all. +@testset "the pullback of a batch allocates linearly in its size ($form)" for form in (:array, :qp) + ps = NetworkParameters((L1 = (W = [0.5 -1 0.25; 1.5 0.75 -2], b = [0.1, -0.3]),)) + model = Chain(ParametricAffine{2}(1)) + function batch(k) + rng = Random.Xoshiro(k) + X, Y = randn(rng, 2, 1, k), randn(rng, 2, 1, k) + μs = [(a = randn(rng),) for _ in 1:k] + form === :array ? (X, Y, μs) : + ((q = X[1:1, :, :], p = X[2:2, :, :]), (q = Y[1:1, :, :], p = Y[2:2, :, :]), μs) + end + b64 = pullback_bytes(model, ps, batch(64)) + b1024 = pullback_bytes(model, ps, batch(1024)) + @test (b1024 - b64) / 960 ≤ 2 * b64 / 64 +end + +# The gradient of the loss with respect to the input, against central finite differences of the same +# loss over the entries of the input. +@testset "the input gradient of ParametricLoss ($form)" for form in (:array, :qp) + ps = NetworkParameters((L1 = (W = [0.5 -1 0.25; 1.5 0.75 -2], b = [0.1, -0.3]),)) + model = Chain(ParametricAffine{2}(1)) + rng = Random.Xoshiro(3) + k = 4 + X, Y = randn(rng, 2, 1, k), randn(rng, 2, 1, k) + μs = [(a = randn(rng),) for _ in 1:k] + as_input(v) = form === :array ? reshape(v, 2, 1, k) : + (q = reshape(v[1:k], 1, 1, k), p = reshape(v[(k + 1):end], 1, 1, k)) + as_vector(x) = form === :array ? vec(x) : vcat(vec(x.q), vec(x.p)) + y = form === :array ? Y : (q = Y[1:1, :, :], p = Y[2:2, :, :]) + loss(x) = ParametricLoss()(model, ps, x, y, μs) + v = form === :array ? vec(X) : vcat(vec(X[1, 1, :]), vec(X[2, 1, :])) + + computed = only(Zygote.gradient(loss, as_input(v))) + @test computed !== nothing + reference = only(grad(central_fdm(5, 1), loss ∘ as_input, v)) + @test norm(as_vector(computed) - reference) ≤ 1e-6 * norm(reference) +end + @testset "the parametric names are public and not exported" begin for name in (:apply_parametric, :ParametricLoss, :ParametricDataLoader) @test Base.ispublic(GeometricMachineLearning, name) From db9bedfb04fb36451e706cd6afc7d909ccc5fdd7 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Sun, 25 Oct 2026 22:20:22 +0100 Subject: [PATCH 4/7] Add K23: the batched ZygotePullback gradient infers as a Union Co-Authored-By: Claude Opus 5.5 --- KNOWN_ISSUES.md | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index 7219b932d..0c7c97b51 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -132,6 +132,22 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. `ParametricDataLoader`'s `(q, p)` constructor has the same one-type restriction and widens with it. +### K23 · `ZygotePullback`'s backward pass on a 3-tensor batch infers as `Union{Nothing, NetworkParameters}`; one sample infers concretely. + +- **location:** `src/pullbacks/zygote_pullback.jl:37-43` +- **kind:** defect +- **found:** 2026-10-25, by the review of the parametric core of #207; the same inference shows on + `origin/main` at `a664f5c9` +- **evidence:** + On Julia 1.13.2, `Base.return_types` of `only(pb(one(v)))` gives + `Union{Nothing, NetworkParameters}` for a batch in each of these forms: `apply_parametric`'s + closure over the columns, a callable struct in place of the closure, the form that indexes one + sample per step, and the user path `ZygotePullback(ParametricLoss())`. One sample gives + `NetworkParameters{Float64, (:L1,), …}`. `main`'s own `ZygotePullback(FeedForwardLoss())` on + `Chain(Dense(2, 2, identity))` and a `2 × 1 × 8` tensor gives + `Any[Union{Nothing, NetworkParameters}]`. The callable struct and the closure give gradients that + differ by `0.0` in `Float32` and in `Float64`. + ## Upstream ### K2 · Revise prints EMFILE errors in the test log From 1f6670cf5f55a6d1966f4da60ed1aadb4af352aa Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Sun, 25 Oct 2026 23:15:00 +0100 Subject: [PATCH 5/7] Renumber the new known issues K22 and K23 to K28 and K29 The CHANGELOG records K22 and K23 as closed, so the two new entries take the next identifiers after the highest one in use, K27. Co-Authored-By: Claude Opus 5.5 --- KNOWN_ISSUES.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index 0c7c97b51..6e54a83e8 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -132,7 +132,7 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. `ParametricDataLoader`'s `(q, p)` constructor has the same one-type restriction and widens with it. -### K23 · `ZygotePullback`'s backward pass on a 3-tensor batch infers as `Union{Nothing, NetworkParameters}`; one sample infers concretely. +### K29 · `ZygotePullback`'s backward pass on a 3-tensor batch infers as `Union{Nothing, NetworkParameters}`; one sample infers concretely. - **location:** `src/pullbacks/zygote_pullback.jl:37-43` - **kind:** defect @@ -163,7 +163,7 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. The blocks are not test failures: both runs are green, and every testset of ea0c397 has the same count in both. -### K22 · AbstractNeuralNetworks' loss norm of a `(q, p)` named tuple promotes `Float32` to `Float64`. +### K28 · AbstractNeuralNetworks' loss norm of a `(q, p)` named tuple promotes `Float32` to `Float64`. - **location:** AbstractNeuralNetworks 0.9.3, `src/losses.jl:73` - **kind:** upstream From c51e94e3fca8aa17be1748f16e6106ac077c6c01 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Mon, 26 Oct 2026 09:55:33 +0100 Subject: [PATCH 6/7] Correct K29: the branch's batch path causes the Union, not main On a664f5c9 and on this branch, the FeedForwardLoss batch gradient infers as a concrete NetworkParameters in Float32 and Float64; only the ParametricLoss batch path gives Union{Nothing, NetworkParameters}. Also match the whole phrase in three ArgumentError message tests instead of a single digit. Co-Authored-By: Claude Opus 5.5 --- KNOWN_ISSUES.md | 17 +++++++++-------- test/data_loader/parametric_data_loader.jl | 8 ++++---- 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index 6e54a83e8..da75e7b00 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -132,21 +132,22 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. `ParametricDataLoader`'s `(q, p)` constructor has the same one-type restriction and widens with it. -### K29 · `ZygotePullback`'s backward pass on a 3-tensor batch infers as `Union{Nothing, NetworkParameters}`; one sample infers concretely. +### K29 · `ZygotePullback(ParametricLoss())`'s backward pass on a 3-tensor batch infers as `Union{Nothing, NetworkParameters}`; one sample infers concretely. -- **location:** `src/pullbacks/zygote_pullback.jl:37-43` +- **location:** `src/parametric.jl:87-101` - **kind:** defect -- **found:** 2026-10-25, by the review of the parametric core of #207; the same inference shows on - `origin/main` at `a664f5c9` +- **found:** #357, review of 2026-10-26; the batch methods of `apply_parametric` cause it, and + `origin/main` at `a664f5c9` does not show it - **evidence:** On Julia 1.13.2, `Base.return_types` of `only(pb(one(v)))` gives `Union{Nothing, NetworkParameters}` for a batch in each of these forms: `apply_parametric`'s closure over the columns, a callable struct in place of the closure, the form that indexes one sample per step, and the user path `ZygotePullback(ParametricLoss())`. One sample gives - `NetworkParameters{Float64, (:L1,), …}`. `main`'s own `ZygotePullback(FeedForwardLoss())` on - `Chain(Dense(2, 2, identity))` and a `2 × 1 × 8` tensor gives - `Any[Union{Nothing, NetworkParameters}]`. The callable struct and the closure give gradients that - differ by `0.0` in `Float32` and in `Float64`. + `NetworkParameters{Float64, (:L1,), …}`. `ZygotePullback(FeedForwardLoss())` on + `Chain(Dense(2, 2, identity))` and a `2 × 1 × 8` tensor gives a concrete + `NetworkParameters{T, (:L1,), …}` in `Float32` and `Float64`, on `a664f5c9` and on this branch. + The callable struct and the closure give gradients that differ by `0.0` in `Float32` and in + `Float64`. ## Upstream diff --git a/test/data_loader/parametric_data_loader.jl b/test/data_loader/parametric_data_loader.jl index ffbfe74e7..9557d0b30 100644 --- a/test/data_loader/parametric_data_loader.jl +++ b/test/data_loader/parametric_data_loader.jl @@ -110,18 +110,18 @@ end # one time step per sample: a sequence length other than one is an error before any step message = argument_error_message(() -> Batch(2, 3, 1)(dl)) - @test occursin("3", message) - @test occursin("sequence length", message) + @test occursin("sequence length of 3", message) # a prediction window other than one is an error before any step, with the time steps to spare message = argument_error_message(() -> Batch{:FeedForward}(2, 1, 2)(dl)) @test occursin("prediction window of 2", message) # a parameter vector of the wrong length, and no trajectory message = argument_error_message(() -> ParametricDataLoader(data, μs[1:2])) - @test occursin("2", message) && occursin("3", message) + @test occursin("hold 3 trajectories", message) && + occursin("there are 2 system parameters", message) message = argument_error_message( () -> ParametricDataLoader(zeros(2d, nt, 0), NamedTuple{(:a,), Tuple{Float64}}[])) - @test occursin("0", message) + @test occursin("hold 0 trajectories", message) # too few time steps for the batch message = argument_error_message(() -> Batch(2)(ParametricDataLoader(data[:, 1:1, :], μs))) From c5409c59c2c14516acbec7e1f9291909b90ddeb3 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Mon, 26 Oct 2026 10:33:58 +0100 Subject: [PATCH 7/7] Infer the batch gradient of ParametricLoss as a concrete NetworkParameters Zygote's adjoint of map adds the gradients of the mapped closure, which holds ps, from nothing, so the gradient of ps on a batch inferred as Union{Nothing, NetworkParameters}. The batch methods of apply_parametric now map their samples through map_samples, whose ChainRulesCore rule pulls each sample back with Zygote.pullback and adds the gradients of ps, as NetworkParameters via map_cotangent, from the first sample on. A (q, p) sample is stacked into one vector inside the map, so that its cotangent reaches Zygote as an array. Outputs and gradients are bitwise unchanged. A test asserts @inferred on the pullback of ParametricLoss for one sample and for an array batch, in Float32 and Float64. K29 leaves KNOWN_ISSUES.md. K30 records the upstream Union of AbstractNeuralNetworks' loss on (q, p) data, which main shows too. Co-Authored-By: Claude Opus 5.5 --- KNOWN_ISSUES.md | 41 +++++++++++++++++++-------------- src/GeometricMachineLearning.jl | 6 +++-- src/parametric.jl | 36 +++++++++++++++++++++++++---- test/parametric.jl | 16 +++++++++++++ 4 files changed, 75 insertions(+), 24 deletions(-) diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index da75e7b00..8bf3625ab 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -132,23 +132,6 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. `ParametricDataLoader`'s `(q, p)` constructor has the same one-type restriction and widens with it. -### K29 · `ZygotePullback(ParametricLoss())`'s backward pass on a 3-tensor batch infers as `Union{Nothing, NetworkParameters}`; one sample infers concretely. - -- **location:** `src/parametric.jl:87-101` -- **kind:** defect -- **found:** #357, review of 2026-10-26; the batch methods of `apply_parametric` cause it, and - `origin/main` at `a664f5c9` does not show it -- **evidence:** - On Julia 1.13.2, `Base.return_types` of `only(pb(one(v)))` gives - `Union{Nothing, NetworkParameters}` for a batch in each of these forms: `apply_parametric`'s - closure over the columns, a callable struct in place of the closure, the form that indexes one - sample per step, and the user path `ZygotePullback(ParametricLoss())`. One sample gives - `NetworkParameters{Float64, (:L1,), …}`. `ZygotePullback(FeedForwardLoss())` on - `Chain(Dense(2, 2, identity))` and a `2 × 1 × 8` tensor gives a concrete - `NetworkParameters{T, (:L1,), …}` in `Float32` and `Float64`, on `a664f5c9` and on this branch. - The callable struct and the closure give gradients that differ by `0.0` in `Float32` and in - `Float64`. - ## Upstream ### K2 · Revise prints EMFILE errors in the test log @@ -181,6 +164,30 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. typeof(AbstractNeuralNetworks._norm(x)) # Float64 ``` +### K30 · The gradient of AbstractNeuralNetworks' loss of a `(q, p)` named tuple infers as `Union{Nothing, NetworkParameters}`. + +- **location:** AbstractNeuralNetworks 0.9.3, `src/losses.jl:86` +- **kind:** upstream +- **found:** #357, review of 2026-10-26; the same inference shows on `origin/main` at `a664f5c9` +- **evidence:** + On Julia 1.13.2 and `a664f5c9`, `Base.return_types` of the pullback of + `ps -> _compute_loss((q = ps.L1.W * q, p = ps.L1.b .* p), (q = yq, p = yp))` gives + `Tuple{Union{Nothing, NetworkParameters}}` in `Float32` and `Float64`, and the same function + with `vcat` in place of the named tuples gives the concrete `NetworkParameters{T, (:L1,), …}`. So + `ZygotePullback(ParametricLoss())` on `(q, p)` data infers its gradient as that `Union`, for one + sample and for a batch, and a plain sum of squares over the `(q, p)` output of a batch infers + concretely. The run-time gradient is a `NetworkParameters`. + + ```julia + using GeometricMachineLearning, Zygote + import AbstractNeuralNetworks + ps = NetworkParameters((L1 = (W = rand(1, 1), b = rand(1)),)) + q, p, yq, yp = rand(1), rand(1), rand(1), rand(1) + v, pb = Zygote.pullback(ps -> AbstractNeuralNetworks._compute_loss( + (q = ps.L1.W * q, p = ps.L1.b .* p), (q = yq, p = yp)), ps) + Base.return_types(pb, (typeof(v),)) # Tuple{Union{Nothing, NetworkParameters}} + ``` + ### K4 · Metal.jl 1.11.1 needs a Metal device to precompile, so a sandboxed run with no Metal cache fails before the tests - **location:** `test/Project.toml:16` diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index 9d61e4901..45532f68e 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -13,8 +13,10 @@ import NeuralNetworkParameters: NetworkParameters using NeuralNetworkParameters: mapstorage # `flatten` puts the parameters of a system into the input of a network that takes them # (`src/parametric.jl`); `parameter_eltype` checks their element type against the data's, and -# `parameterlayout` their shape against each other's. -using NeuralNetworkParameters: flatten, parameter_eltype, parameterlayout +# `parameterlayout` their shape against each other's. `map_cotangent` and `storage_gradient` turn the +# cotangent of a parameter set into its gradient in the reverse rule of a batch. +using NeuralNetworkParameters: flatten, parameter_eltype, parameterlayout, map_cotangent, + storage_gradient using ChainRulesCore using GeometricBase using GeometricSolutions: GeometricSolution, EnsembleSolution, TimeSeries diff --git a/src/parametric.jl b/src/parametric.jl index 5f3a6b02d..7767ce637 100644 --- a/src/parametric.jl +++ b/src/parametric.jl @@ -87,22 +87,48 @@ end function apply_parametric(model::Chain, x::AbstractArray{<:Any, 3}, μs::AbstractVector, ps::NetworkParameters) check_parametric_batch(size(x), μs) - outputs = map((xᵢ, μ) -> apply_parametric(model, xᵢ, μ, ps), sample_columns(x), μs) + outputs = map_samples((xᵢ, μ, ps) -> apply_parametric(model, xᵢ, μ, ps), ps, + sample_columns(x), μs) reshape(reduce(hcat, outputs), :, 1, length(μs)) end +# The output of a sample is stacked, `q` above `p`, so that its cotangent is an array, which the rule +# of `map_samples` hands to Zygote as it is. function apply_parametric(model::Chain, x::NamedTupleOfArrays{<:Any, (:q, :p)}, μs::AbstractVector, ps::NetworkParameters) check_parametric_batch(size(x.q), μs) - outputs = map((q, p, μ) -> apply_parametric(model, (q = q, p = p), μ, ps), - sample_columns(x.q), sample_columns(x.p), μs) - (q = reshape(reduce(hcat, [output.q for output in outputs]), :, 1, length(μs)), - p = reshape(reduce(hcat, [output.p for output in outputs]), :, 1, length(μs))) + outputs = map_samples(ps, sample_columns(x.q), sample_columns(x.p), μs) do q, p, μ, ps + output = apply_parametric(model, (q = q, p = p), μ, ps) + vcat(output.q, output.p) + end + stacked = reduce(hcat, outputs) + d = size(x.q, 1) + (q = reshape(stacked[1:d, :], :, 1, length(μs)), + p = reshape(stacked[(d + 1):end, :], :, 1, length(μs))) end # The samples of a `d × 1 × k` batch, as the `k` columns of its `d × k` reshape. sample_columns(x::AbstractArray{<:Any, 3}) = eachcol(reshape(x, size(x, 1), size(x, 3))) +# `f(sample..., ps)` for each sample, where sample `i` is element `i` of each of `samples`. +map_samples(f, ps, samples...) = map((sample...) -> f(sample..., ps), samples...) + +# Zygote's adjoint of `map` adds the gradients of the closure it maps, which holds `ps`, from `nothing`, +# so the gradient of `ps` infers as `Union{Nothing, NetworkParameters}`. This rule pulls each sample +# back on its own and adds the gradients of `ps`, as `NetworkParameters`, from the first sample on; a +# batch holds at least one sample. `f` holds the `Chain`, whose trainable parameters are in `ps`, so +# `f` has no gradient. +function ChainRulesCore.rrule(::typeof(map_samples), f, ps::NetworkParameters, samples...) + outputs_and_pullbacks = map((sample...) -> Zygote.pullback(f, sample..., ps), samples...) + function map_samples_pullback(Δoutputs) + Δs = map(((_, pullback), Δ) -> pullback(Δ), outputs_and_pullbacks, unthunk(Δoutputs)) + Δps = sum(Δ -> map_cotangent(storage_gradient, ps, last(Δ)), Δs) + Δsamples = ntuple(i -> map(Δ -> Δ[i], Δs), length(samples)) + (NoTangent(), NoTangent(), Δps, Δsamples...) + end + map(first, outputs_and_pullbacks), map_samples_pullback +end + # A batch holds one time step per sample and one `μ` per sample. function check_parametric_batch(input_size::NTuple{3, Int}, μs::AbstractVector) input_size[2] == 1 && input_size[3] == length(μs) || throw(DimensionMismatch( diff --git a/test/parametric.jl b/test/parametric.jl index aa54b238d..77d5ce3fa 100644 --- a/test/parametric.jl +++ b/test/parametric.jl @@ -152,6 +152,22 @@ end @test norm(as_vector(computed) - reference) ≤ 1e-6 * norm(reference) end +# The gradient that the pullback of `ParametricLoss` returns infers as its run-time type, for one sample +# and for a batch. The `(q, p)` form is left out: AbstractNeuralNetworks' loss of a `(q, p)` named tuple +# infers its gradient as a `Union` on its own (KNOWN_ISSUES K30). +@testset "the gradient of ParametricLoss infers concretely ($T)" for T in REAL_ELTYPES + ps = NetworkParameters((L1 = (W = T[0.5 -1 0.25; 1.5 0.75 -2], b = T[0.1, -0.3]),)) + model = Chain(ParametricAffine{2}(1)) + rng = Random.Xoshiro(5) + k = 4 + X, Y = randn(rng, T, 2, 1, k), randn(rng, T, 2, 1, k) + μs = [(a = randn(rng, T),) for _ in 1:k] + for data in ((X[:, 1, 1], Y[:, 1, 1], μs[1]), (X, Y, μs)) + value, back = ZygotePullback(ParametricLoss())(ps, model, data) + @test (@inferred back(one(value))) isa NetworkParameters{T} + end +end + @testset "the parametric names are public and not exported" begin for name in (:apply_parametric, :ParametricLoss, :ParametricDataLoader) @test Base.ispublic(GeometricMachineLearning, name)