From 98c1c76495b2a02a30cb20f71887de9389fc54ce Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 16 Oct 2026 08:00:42 +0200 Subject: [PATCH 1/8] DataLoader: samples against trajectories, one count and one gather (part M2) A loader built from an input and an output holds a set of samples whatever its kind, so batching no longer branches on the kind alone: the target is the output itself, never the input shifted by the sequence length. One count, nn = the start positions of one parameter index, replaces the two the batching had, and it draws every sample of a supervised loader including the last. The five gather methods and four kernels become one kernel and one helper; the copy loops of a solution and an ensemble become one `data_tensor`; `map_to_*` reaches the two tensors of a `(q, p)` pair with `map`; `tensor_assign.jl` keeps only the crop of the output. autoencoder defaults to false in every constructor, and `DataLoader(input, output)` takes two free element types, so a real input with an integer one-hot target is one call. The two output fields of the struct carry a fifth type parameter, which makes them concrete. KNOWN_ISSUES B10, C17, K8 and K12 leave the file: the ensemble method is reachable through a rigid-body `odeensemble`, and it has a test. --- KNOWN_ISSUES.md | 79 ----- docs/src/data_loader/data_loader.md | 12 +- .../symplectic_autoencoders/training.jl | 2 +- src/data_loader/batch.jl | 313 +++++----------- src/data_loader/data_loader.jl | 333 +++++++----------- src/data_loader/tensor_assign.jl | 66 +--- test/data_loader/accuracy.jl | 42 ++- test/data_loader/autoencoder_keyword.jl | 5 +- test/data_loader/batch.jl | 169 +++++++++ test/data_loader/batch_data_loader_qp_test.jl | 5 +- test/data_loader/batch_functor.jl | 2 +- test/data_loader/batch_index_set.jl | 2 +- test/data_loader/data_loader.jl | 162 +++++++++ .../optimizer_functor_with_adam.jl | 6 +- test/integration/docstrings/data_loader.jl | 11 +- test/kernels/kernel_pullbacks.jl | 2 - test/quality/jet.jl | 54 ++- test/runtests.jl | 2 + 18 files changed, 652 insertions(+), 615 deletions(-) create mode 100644 test/data_loader/batch.jl create mode 100644 test/data_loader/data_loader.jl diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index eb936965a..6158d6831 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -4,43 +4,6 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. ## B. Known defects -### 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` -- **kind:** missing test -- **found:** 2026-09-18 -- **evidence:** - It dispatches - on `ST <: Union{GeometricSolution{T, T1, TT, NamedTuple{(:t, :q, :v), TuT}}, - GeometricSolution{T, T1, TT, NamedTuple{(:t, :q, :q̇), TuT}}}` — a `GeometricSolution` whose - `dataser` has exactly the two keys `:q` and `:v` (or `:q̇`) besides `:t`, with no `:p`. - - Checked against `GeometricEquations` 0.21.3 (this package's resolved version): every equation - type's own `initialstate(equ, t, ics, params)` reconstructs its `ics` from its own fixed field - set regardless of what is passed in, and no type pairs `:v`/`:q̇` without also carrying `:p` — - `SODE`/`ODE` give `(:q,)` alone; `PODE`/`HODE` give `(:q, :p)`; `IODE`/`LODE` give - `(:q, :p, :v)`; `IDAE`/`LDAE` give `(:q, :p, :v, :λ, :μ)`. Verified directly for `SODE`: - `initialstate(equ::SODE, t, ics, params) = (q = _statevariable(ics.q, periodicity(equ)),)` - discards everything but `.q` even when `ics` already has a `:v` key. So no - `EquationProblem`/`EnsembleProblem` built from any equation type this package depends on can - produce a two-key `dataser` — passing a NamedTuple with the right keys through the public - constructor does not help, because the equation-specific `initialstate` method throws it away. - - `GeometricSolution` and `EnsembleSolution` each define exactly one inner constructor (taking a - `GeometricProblem`/`EnsembleProblem`), so Julia generates no default all-fields constructor for - either, and there is no supported way to build one directly. The two low-level bypasses tried — - `ccall(:jl_new_struct, ...)` on the (mutable) `GeometricSolution`, and - `ccall(:jl_new_struct_uninit, ...)` followed by `setfield!` on each field — both crashed the - Julia process with a segmentation fault (the first immediately; the second on a later, unrelated - allocation, consistent with GC scanning a partially-initialized object), which is why neither is - a technique this package's test suite should rely on. - - The fix in `src/data_loader/data_loader.jl:339` (`zeros(T, ...)` instead of untyped `zeros(...)`) - is correct by inspection — it matches the file's three sibling constructors character for - character — but is untested. Closing this means either GeometricEquations gaining an equation - type whose `initialstate` returns exactly `(:q, :v)` or `(:q, :q̇)`, or deciding the method is - dead code and removing it (Part E of the audit, not this one). - ### B11 · `PositionalEncoding` is CPU-only, and the failure is a compile error rather than a slowdown. - **location:** — @@ -126,23 +89,6 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. deliberately, and the resolution is upstream declaring them `public`. Each check's reason is written beside it in the test file. -### C17 · `accuracy` cannot be reached through any public `DataLoader` constructor. - -- **location:** `src/data_loader/data_loader.jl:487` -- **kind:** defect -- **found:** 2026-09-18 -- **evidence:** - Its signature is - `DataLoader{T, AT <: AbstractArray{T}, BT <: AbstractArray{T1}}` with `T1 <: Integer` - (`src/data_loader/data_loader.jl:487`), and every `DataLoader(input, output)` method requires the - two arrays to share one element type. A classifier's data are the case where they do not: real - inputs, integer one-hot targets. So the only way to build an argument for this function is to name - the parametric type, which is what `test/data_loader/accuracy.jl` does. - - Closing it is a `DataLoader(input::AbstractArray{T, 3}, output::AbstractArray{T1, 3})` constructor - with the two element types free. That is an API addition rather than a cleanup, which is why the - test records the gap instead of the same change closing it. - ### C18 · `src/kernels/exponentials/tensor_exponential.jl` no longer defines `tensor_exponential`; it holds the identity tensor and its `rrule`, used by `tensor_cayley.jl` and `cpu_inverse.jl`. - **location:** `src/kernels/exponentials/tensor_exponential.jl` @@ -231,16 +177,6 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. `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. -### K8 · A `DataLoader` of one sample cannot be batched. - -- **location:** `src/data_loader/batch.jl` -- **kind:** found late -- **found:** 2026-10-12, by the critics of part M1a -- **evidence:** - `DataLoader(x, x)` with `x = rand(Float32, 2, 1)` and `Batch(1)` raises the time-series - assertion "The number of time steps has to be greater than sequence length + prediction window" - when `train!` draws its minibatches. - ### K11 · A comment in `adam_with_learning_rate_decay.jl` reports a past experiment. - **location:** `test/optimizers/adam_with_learning_rate_decay.jl` @@ -251,21 +187,6 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. on 1.13 and *rose* on 1.12". The fix states only the present constraint: a 32-epoch budget drives the learning rate to `η₂` almost at once, and 128 epochs train. -### K12 · `Batch` never draws the last sample of a `DataLoader` of matrices. - -- **location:** `src/data_loader/batch.jl` -- **kind:** found late -- **found:** 2026-10-12, by the critics of part M1a -- **evidence:** - On a `DataLoader(x, x)` with `x = rand(2, 10)`, the minibatches of `Batch(bs)` hold the samples - 1 to 9 for each `bs` in 1, 3, 4 and 10; sample 10 is in none of them. `Batch(1)` gives 9 - minibatches. - - ``` - 1 nbatches 9 samples [(1, 1), (2, 1), (3, 1), (4, 1), (5, 1), (6, 1), (7, 1), (8, 1), (9, 1)] - 10 nbatches 1 samples [(1, 1), (2, 1), (3, 1), (4, 1), (5, 1), (6, 1), (7, 1), (8, 1), (9, 1)] - ``` - ### K15 · `HNNLoss` and `SymbolicPullback(arch)` accept `(q, p)` data in their signatures and fail on it inside the generated vector field. - **location:** `src/loss/hnn_loss.jl:33` diff --git a/docs/src/data_loader/data_loader.md b/docs/src/data_loader/data_loader.md index 22b34f1e6..58b8f72a1 100644 --- a/docs/src/data_loader/data_loader.md +++ b/docs/src/data_loader/data_loader.md @@ -8,7 +8,7 @@ The data loader can be called with various types of arrays as input, for example using GeometricMachineLearning # hide SnapshotMatrix = [ 1; 2;; 3; 4;; 5; 6;; 7; 8;; 9; 10 ] -dl = DataLoader(SnapshotMatrix; suppress_info = true) +dl = DataLoader(SnapshotMatrix; autoencoder = true, suppress_info = true) ``` or a snapshot tensor: @@ -20,13 +20,13 @@ SnapshotTensor = [ 1; 2;; 3; 4;; 5; 6;; 7; 8;; 9; 10 ;;;] dl = DataLoader(SnapshotTensor; suppress_info = true) ``` -Here the `DataLoader` has different properties `:RegularData` and `:TimeSeries`. This indicates that in the first case we treat all columns in the input tensor independently; this is mostly used for [autoencoder problems](@ref "Autoencoders"). In the second case we have *time series-like data*, which are mostly used for [integration problems](@ref "Neural Network Integrators"). As shown above the default when using a matrix is `:RegularData` and the default when using a tensor is `:TimeSeries`. +Here the `DataLoader` has different properties `:RegularData` and `:TimeSeries`. This indicates that in the first case we treat all columns in the input tensor independently; this is mostly used for [autoencoder problems](@ref "Autoencoders"). In the second case we have *time series-like data*, which are mostly used for [integration problems](@ref "Neural Network Integrators"). The default is `:TimeSeries` for a matrix and for a tensor alike, and `autoencoder = true` is what makes a matrix a set of samples; a loader that is given an *output* as well, `DataLoader(input, output)`, always holds a set of samples, whatever its kind. -We can also treat a problem with a matrix as input as a time series-like problem by providing an additional keyword argument: `autoencoder=false`: +`autoencoder = false`, the default, leaves a matrix as a time series: ```@example snapshot_matrix dl = DataLoader(SnapshotMatrix; autoencoder=false, suppress_info = true) -@assert typeof(dl) == DataLoader{Int64, Array{Int64, 3}, Nothing, :TimeSeries} # hide +@assert typeof(dl) == DataLoader{Int64, Array{Int64, 3}, Nothing, :TimeSeries, Nothing} # hide dl |> typeof ``` @@ -36,7 +36,7 @@ If we deal with hamiltonian systems we typically split the coordinates into a `` using GeometricMachineLearning # hide SymplecticSnapshotTensor = (q = SnapshotTensor, p = SnapshotTensor) dl = DataLoader(SymplecticSnapshotTensor) -@assert typeof(dl) == DataLoader{Int64, @NamedTuple{q::Array{Int64, 3}, p::Array{Int64, 3}}, Nothing, :TimeSeries} # hide +@assert typeof(dl) == DataLoader{Int64, @NamedTuple{q::Array{Int64, 3}, p::Array{Int64, 3}}, Nothing, :TimeSeries, Nothing} # hide dl |> typeof ``` @@ -106,7 +106,7 @@ batch = Batch(3) batch(dl) ``` -In those two examples the `autoencoder` keyword was set to `true` (the default). This is why the first index was always `1`. This changes if we set `autoencoder = false`: +In the first of those two examples the `autoencoder` keyword was set to `true`, which is why the first index was always `1`. The second sets it to `false`, the default: ```@example using GeometricMachineLearning # hide diff --git a/scripts/reproduction/symplectic_autoencoders/training.jl b/scripts/reproduction/symplectic_autoencoders/training.jl index b0fed708f..e32542456 100644 --- a/scripts/reproduction/symplectic_autoencoders/training.jl +++ b/scripts/reproduction/symplectic_autoencoders/training.jl @@ -68,7 +68,7 @@ catch cpu_backend() end -dl = DataLoader(data) +dl = DataLoader(data; autoencoder = true) n_time_steps = size(data, 2) / n_params N = size(data, 1)÷2 diff --git a/src/data_loader/batch.jl b/src/data_loader/batch.jl index 16c3c9835..a9113c511 100644 --- a/src/data_loader/batch.jl +++ b/src/data_loader/batch.jl @@ -22,7 +22,7 @@ import Random rng = Random.TaskLocalRNG() Random.seed!(rng, 123) -dl = DataLoader(rand(rng, 5)) +dl = DataLoader(rand(rng, 5); autoencoder = true) batch = Batch(2) batches = batch(dl) @@ -98,14 +98,11 @@ function Batch(batch_size::Int, seq_length::Int, prediction_window::Int = seq_le Batch{:Transformer}(batch_size, seq_length, prediction_window) end -function Batch(::Int, ::Nothing, ::Int) - error("Cannot provide prediction window alone. Need sequence length!") -end - +# The classifier's batch: one draw per sample, with the whole input sequence and the whole output. +# Both lengths are the loader's. Its old prediction window of `0` would gather an output of width 0. function Batch(batch_size::Int, - dl::DataLoader{T, AT, OT, :TimeSeries}) where { - T, T1, AT <: AbstractArray{T, 3}, OT <: AbstractArray{T1, 3}} - Batch(batch_size, dl.input_time_steps, 0) + dl::DataLoader{T, AT, OT, K, Int}) where {T, AT, OT <: AbstractArray, K} + Batch(batch_size, dl.input_time_steps, dl.output_time_steps) end @doc raw""" @@ -113,7 +110,10 @@ end Compute the number of batches. -Here the distinction is between data that are *time-series like* and data that are *autoencoder like*. +One epoch draws every start position of every parameter index once, in minibatches of the batch +size, so the count is `cld(n · n_params, batch_size)` with `n` the number of those positions and +`samples without replacement`: two indices are never drawn twice. The last minibatch holds what is +left. # Examples @@ -148,99 +148,91 @@ true Here we see that in the *autoencoder case* that last minibatch has an additional element. """ -function number_of_batches(dl::DataLoader{T, AT, OT, :TimeSeries}, - batch::Batch) where {T, BT <: AbstractArray{T, 3}, - AT <: Union{BT, NamedTuple{(:q, :p), Tuple{BT, BT}}}, OT} - @assert dl.input_time_steps ≥ (batch.seq_length + batch.prediction_window) "The number of time steps has to be greater than sequence length + prediction window." - Int(ceil((dl.input_time_steps - (batch.seq_length - 1) - batch.prediction_window) * - dl.n_params / batch.batch_size)) +number_of_batches(dl::DataLoader, batch::Batch) = cld( + number_of_start_positions(dl, batch) * dl.n_params, batch.batch_size) + +# The start positions of one parameter index: the input window must fit the input's time axis, and +# the target must fit its own. The target is the output at offset `0` where the loader holds one, +# the input at offset `seq_length` where it is a time series without an output, and there is none +# where the loader is a set of samples without an output. +function number_of_start_positions(dl::DataLoader{T, AT, OT, K}, + batch::Batch) where {T, AT, OT, K} + input_window = dl.input_time_steps - batch.seq_length + 1 + target_window = if OT <: AbstractArray + dl.output_time_steps - batch.prediction_window + 1 + elseif K === :TimeSeries + input_window - batch.prediction_window + else + input_window + end + check_window(min(input_window, target_window), dl, batch) end -function number_of_batches(dl::DataLoader{T, AT, OT, :RegularData}, - batch::Batch) where {T, BT <: AbstractArray{T, 3}, - AT <: Union{BT, NamedTuple{(:q, :p), Tuple{BT, BT}}}, OT} - Int(ceil(dl.input_time_steps * dl.n_params / batch.batch_size)) +# A loader and a batch that leave no start position are a caller's mistake: an epoch of no minibatch +# trains nothing and says nothing, so the three numbers that decide it are named here. +function check_window(n::Int, dl::DataLoader, batch::Batch) + n ≥ 1 || + throw(ArgumentError("the `DataLoader` has $(dl.input_time_steps) time steps, and the " * + "`Batch` asks for a sequence length of $(batch.seq_length) and a " * + "prediction window of $(batch.prediction_window), which leaves no " * + "start position.")) + n end -function batch_over_two_axes(batch::Batch, number_columns::Int, third_dim::Int, dl::DataLoader) - time_indices = shuffle(1:number_columns) - parameter_indices = shuffle(1:third_dim) - complete_indices = Iterators.product(time_indices, parameter_indices) |> collect |> vec - n_batches = number_of_batches(dl, batch) - batches = Vector{Vector{Tuple{Int, Int}}}(undef, n_batches) - for batch_number in 1:(n_batches - 1) - batches[batch_number] = complete_indices[((batch_number - 1) * batch.batch_size + 1):(batch_number * batch.batch_size)] +# The index set of one epoch: the product of the shuffled start positions and parameter indices, cut +# into minibatches of the batch size. +function batch_over_two_axes(batch::Batch, number_columns::Int, third_dim::Int) + complete_indices = vec( + collect(Iterators.product(shuffle(1:number_columns), shuffle(1:third_dim)))) + batches = Vector{Vector{Tuple{Int, Int}}}() + for i in 1:cld(length(complete_indices), batch.batch_size) + push!(batches, + complete_indices[((i - 1) * batch.batch_size + 1):min(i * batch.batch_size, end)]) end - batches[n_batches] = complete_indices[((n_batches - 1) * batch.batch_size + 1):end] batches end -function (batch::Batch)(dl::DataLoader{T, - BT, - OT, - :RegularData}) where {T, AT <: AbstractArray{T, 3}, - BT <: Union{AT, NamedTuple{(:q, :p), Tuple{AT, AT}}}, OT} - batch_over_two_axes(batch, dl.input_time_steps, dl.n_params, dl) -end - -function (batch::Batch)(dl::DataLoader{T, - BT, - OT, - :TimeSeries}) where {T, AT <: AbstractArray{T, 3}, - BT <: Union{AT, NamedTuple{(:q, :p), Tuple{AT, AT}}}, OT} - batch_over_two_axes( - batch, dl.input_time_steps - (batch.seq_length - 1) - batch.prediction_window, - dl.n_params, dl) -end - -@kernel function assign_input_from_vector_of_tuples_kernel!( - q_input::AT, p_input::AT, input::NamedTuple{(:q, :p), Tuple{AT, AT}}, - indices::AbstractArray{Int, 2}) where {T, AT <: AbstractArray{T, 3}} - i, j, k = @index(Global, NTuple) - - q_input[i, j, k] = input.q[i, indices[1, k] + j - 1, indices[2, k]] - p_input[i, j, k] = input.p[i, indices[1, k] + j - 1, indices[2, k]] +function (batch::Batch)(dl::DataLoader) + batch_over_two_axes(batch, number_of_start_positions(dl, batch), dl.n_params) end -@kernel function assign_input_from_vector_of_tuples_kernel!(input::AT, data::AT, - indices::AbstractArray{Int, 2}) where {T, AT <: AbstractArray{T, 3}} +@kernel function assign_from_vector_of_tuples_kernel!( + input::AT, data::BT, indices::AbstractArray{Int, 2}, + offset) where { + T, AT <: AbstractArray{T, 3}, BT <: AbstractArray{T, 3}} i, j, k = @index(Global, NTuple) - input[i, j, k] = data[i, indices[1, k] + j - 1, indices[2, k]] + input[i, j, k] = data[i, indices[1, k] + offset + j - 1, indices[2, k]] end -@kernel function assign_output_from_vector_of_tuples_kernel!( - q_output::AT, p_output::AT, input::NamedTuple{(:q, :p), Tuple{AT, AT}}, - indices::AbstractArray{Int, 2}, - seq_length::Int) where {T, AT <: AbstractArray{T, 3}} - i, j, k = @index(Global, NTuple) - - q_output[i, j, k] = input.q[i, indices[1, k] + seq_length + j - 1, indices[2, k]] - p_output[i, j, k] = input.p[i, indices[1, k] + seq_length + j - 1, indices[2, k]] +# A kernel cannot index a `Vector{Tuple{Int, Int}}`, so the batch indices go in as a `2 × n` matrix: +# the host builds it and the backend takes it in one copy. +function convert_vector_of_tuples_to_matrix(backend::Backend, + batch_indices_tuple::Vector{Tuple{Int, Int}}) + batch_indices = zeros(Int, 2, length(batch_indices_tuple)) + for t in axes(batch_indices_tuple, 1) + batch_indices[1, t], batch_indices[2, t] = batch_indices_tuple[t] + end + map_to_new_backend(batch_indices, backend) end -@kernel function assign_output_from_vector_of_tuples_kernel!( - output::AT, data::AT, indices::AbstractArray{Int, 2}, - seq_length::Int) where {T, AT <: AbstractArray{T, 3}} - i, j, k = @index(Global, NTuple) - - output[i, j, k] = data[i, indices[1, k] + seq_length + j - 1, indices[2, k]] +# The gather kernel, applied to a tensor or to the `q` and `p` of a `(q, p)` pair. The gathered +# tensor holds `width` positions from `offset` on, for every drawn start position and parameter +# index, and has the element type of the data it reads — so one integer target stays integer. The +# first axis is the loader's own `output_dim`/`input_dim`, halved for each of a `(q, p)` pair. +function gather(data::AbstractArray{T, 3}, dim::Int, indices::AbstractArray{Int, 2}, + offset::Int, width::Int, batch_size::Int) where {T} + backend = networkbackend(data) + gathered = KernelAbstractions.allocate(backend, T, dim, width, batch_size) + gather! = assign_from_vector_of_tuples_kernel!(backend) + gather!(gathered, data, indices, offset, ndrange = size(gathered)) + gathered end -# A kernel cannot index a `Vector{Tuple{Int, Int}}`, so the batch indices go in as a `2 × n` matrix. -function convert_vector_of_tuples_to_matrix(backend::Backend, batch_indices_tuple::Vector{Tuple{ - Int, Int}}) - _batch_size = length(batch_indices_tuple) - - batch_indices = KernelAbstractions.allocate(backend, Int, 2, _batch_size) - batch_indices_temp = zeros(Int, size(batch_indices)...) - for t in axes(batch_indices_tuple, 1) - batch_indices_temp[1, t] = batch_indices_tuple[t][1] - batch_indices_temp[2, t] = batch_indices_tuple[t][2] - end - batch_indices = typeof(batch_indices)(batch_indices_temp) - - batch_indices +function gather(data::NamedTuple{(:q, :p), Tuple{AT, AT}}, dim::Int, + indices::AbstractArray{Int, 2}, offset::Int, width::Int, + batch_size::Int) where {T, AT <: AbstractArray{T, 3}} + map(tensor -> gather(tensor, dim ÷ 2, indices, offset, width, batch_size), data) end @doc raw""" @@ -253,7 +245,8 @@ Assign batch data based on (i) input and (ii) batch indices. ```jldoctest using GeometricMachineLearning -dl = DataLoader([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]; suppress_info = true) +dl = DataLoader([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]; + autoencoder = true, suppress_info = true) batch = Batch(3) batch_indices = [(1, 1), (1, 3), (1, 5)] @@ -272,139 +265,21 @@ GeometricMachineLearning.convert_input_and_batch_indices_to_array(dl, batch, bat 0.5 ``` """ -function convert_input_and_batch_indices_to_array(dl::DataLoader{T, BT}, - batch::Batch, - batch_indices_tuple::Vector{Tuple{Int, Int}}) where { - T, AT <: AbstractArray{T, 3}, BT <: NamedTuple{(:q, :p), Tuple{AT, AT}}} - backend = networkbackend(dl.input.q) - - # the batch size is smaller for the last batch - _batch_size = length(batch_indices_tuple) - - batch_indices = convert_vector_of_tuples_to_matrix(backend, batch_indices_tuple) - - q_input = KernelAbstractions.allocate( - backend, T, dl.input_dim ÷ 2, batch.seq_length, _batch_size) - p_input = similar(q_input) - - assign_input_from_vector_of_tuples! = assign_input_from_vector_of_tuples_kernel!(backend) - assign_input_from_vector_of_tuples!(q_input, p_input, dl.input, batch_indices, - ndrange = (dl.input_dim ÷ 2, batch.seq_length, _batch_size)) - - q_output = KernelAbstractions.allocate( - backend, T, dl.input_dim ÷ 2, batch.prediction_window, _batch_size) - p_output = similar(q_output) - - assign_output_from_vector_of_tuples! = assign_output_from_vector_of_tuples_kernel!(backend) - assign_output_from_vector_of_tuples!( - q_output, p_output, dl.input, batch_indices, batch.seq_length, - ndrange = (dl.input_dim ÷ 2, batch.prediction_window, _batch_size)) - - (q = q_input, p = p_input), (q = q_output, p = p_output) -end - -function convert_input_and_batch_indices_to_array(dl::DataLoader{T, BT}, batch::Batch, - batch_indices_tuple::Vector{Tuple{Int, Int}}) where {T, BT <: AbstractArray{T, 3}} - backend = networkbackend(dl.input) - - # the batch size is smaller for the last batch - _batch_size = length(batch_indices_tuple) - - batch_indices = convert_vector_of_tuples_to_matrix(backend, batch_indices_tuple) - - input = KernelAbstractions.allocate( - backend, T, dl.input_dim, batch.seq_length, _batch_size) - - assign_input_from_vector_of_tuples! = assign_input_from_vector_of_tuples_kernel!(backend) - assign_input_from_vector_of_tuples!(input, dl.input, batch_indices, - ndrange = (dl.input_dim, batch.seq_length, _batch_size)) - - output = KernelAbstractions.allocate( - backend, T, dl.input_dim, batch.prediction_window, _batch_size) - - assign_output_from_vector_of_tuples! = assign_output_from_vector_of_tuples_kernel!(backend) - assign_output_from_vector_of_tuples!(output, dl.input, batch_indices, batch.seq_length, - ndrange = (dl.input_dim, batch.prediction_window, _batch_size)) - - input, output -end - -function convert_input_and_batch_indices_to_array( - dl::DataLoader{T, BT, Nothing, :RegularData}, - batch::Batch, - batch_indices_tuple::Vector{Tuple{Int, Int}}) where { - T, AT <: AbstractArray{T, 3}, BT <: NamedTuple{(:q, :p), Tuple{AT, AT}}} - backend = networkbackend(dl.input.q) - - # the batch size is smaller for the last batch - _batch_size = length(batch_indices_tuple) - - batch_indices = convert_vector_of_tuples_to_matrix(backend, batch_indices_tuple) - - q_input = KernelAbstractions.allocate( - backend, T, dl.input_dim ÷ 2, batch.seq_length, _batch_size) - p_input = similar(q_input) - - assign_input_from_vector_of_tuples! = assign_input_from_vector_of_tuples_kernel!(backend) - assign_input_from_vector_of_tuples!(q_input, p_input, dl.input, batch_indices, - ndrange = (dl.input_dim ÷ 2, batch.seq_length, _batch_size)) - - (q = q_input, p = p_input) -end - function convert_input_and_batch_indices_to_array( - dl::DataLoader{T, BT, Nothing, :RegularData}, batch::Batch, - batch_indices_tuple::Vector{Tuple{Int, Int}}) where {T, BT <: AbstractArray{T, 3}} - backend = networkbackend(dl.input) - - # the batch size is smaller for the last batch - _batch_size = length(batch_indices_tuple) - - batch_indices = convert_vector_of_tuples_to_matrix(backend, batch_indices_tuple) - - input = KernelAbstractions.allocate( - backend, T, dl.input_dim, batch.seq_length, _batch_size) - - assign_input_from_vector_of_tuples! = assign_input_from_vector_of_tuples_kernel!(backend) - assign_input_from_vector_of_tuples!(input, dl.input, batch_indices, - ndrange = (dl.input_dim, batch.seq_length, _batch_size)) - - input -end - -# for the case when the DataLoader also contains an output -function convert_input_and_batch_indices_to_array(dl::DataLoader{T, BT, OT}, - ::Batch, - batch_indices_tuple::Vector{Tuple{Int, Int}}) where { - T, T1, BT <: AbstractArray{T, 3}, OT <: AbstractArray{T1, 3}} - parameter_indices = [batch_index[2] for batch_index in batch_indices_tuple] - @views input_batch = dl.input[:, :, parameter_indices] - @views output_batch = dl.output[:, :, parameter_indices] - - input_batch, output_batch -end - -function convert_input_and_batch_indices_to_array(dl::DataLoader{T, BT, BT}, batch::Batch, - batch_indices_tuple::Vector{Tuple{Int, Int}}) where {T, BT <: AbstractArray{T, 3}} - backend = networkbackend(dl.input) - - # the batch size is smaller for the last batch - _batch_size = length(batch_indices_tuple) - - batch_indices = convert_vector_of_tuples_to_matrix(backend, batch_indices_tuple) - - input = KernelAbstractions.allocate( - backend, T, dl.input_dim, batch.seq_length, _batch_size) - - assign_input_from_vector_of_tuples! = assign_input_from_vector_of_tuples_kernel!(backend) - assign_input_from_vector_of_tuples!(input, dl.input, batch_indices, - ndrange = (dl.input_dim, batch.seq_length, _batch_size)) - - output = KernelAbstractions.allocate( - backend, T, dl.output_dim, batch.prediction_window, _batch_size) - - assign_input_from_vector_of_tuples!(output, dl.output, batch_indices, - ndrange = (dl.output_dim, batch.prediction_window, _batch_size)) - - input, output + dl::DataLoader{T, AT, OT, K}, batch::Batch, + batch_indices_tuple::Vector{Tuple{Int, Int}}) where {T, AT, OT, K} + batch_size = length(batch_indices_tuple) + indices = convert_vector_of_tuples_to_matrix(networkbackend(dl), batch_indices_tuple) + input = gather(dl.input, dl.input_dim, indices, 0, batch.seq_length, batch_size) + # a loader that holds an output is a set of samples whatever its kind: the target is its output + if dl.output !== nothing + return input, + gather(dl.output, dl.output_dim, indices, 0, batch.prediction_window, + batch_size) + end + # without one, a time series predicts its own continuation and a set of samples has no target + K === :TimeSeries || return input + input, + gather(dl.input, dl.input_dim, indices, batch.seq_length, batch.prediction_window, + batch_size) end diff --git a/src/data_loader/data_loader.jl b/src/data_loader/data_loader.jl index b4827fb06..cc56b417f 100644 --- a/src/data_loader/data_loader.jl +++ b/src/data_loader/data_loader.jl @@ -29,30 +29,40 @@ dl.input # output -[ Info: You have provided a matrix as input. The axes will be interpreted as (i) system dimension and (ii) number of parameters. -2×1×3 Array{Int64, 3}: +[ Info: You have provided a matrix as input. The axes will be interpreted as (i) system dimension and (ii) number of time steps. +2×3×1 Array{Int64, 3}: [:, :, 1] = - 1 - 4 - -[:, :, 2] = - 2 - 5 - -[:, :, 3] = - 3 - 6 + 1 2 3 + 4 5 6 ``` """ struct DataLoader{T, AT <: Union{NamedTuple, AbstractArray{T}}, - OT <: Union{AbstractArray, Nothing}, DataType} + OT <: Union{AbstractArray, Nothing}, K, S <: Union{Int, Nothing}} input::AT output::OT input_dim::Int input_time_steps::Int n_params::Int - output_dim::Union{Int, Nothing} - output_time_steps::Union{Int, Nothing} + output_dim::S + output_time_steps::S +end + +# The fifth parameter is the type of the two output fields, which is `Nothing` for a loader without +# an output and `Int` with one. A caller that names the first four — `DataLoader{T, AT, OT, +# :TimeSeries}(…)`, which is how `GMLDatasets` builds a loader — gets it filled in from the fields. +function DataLoader{T, AT, OT, K}(input::AT, output::OT, input_dim::Int, + input_time_steps::Int, n_params::Int, output_dim, output_time_steps) where { + T, AT, OT, K} + DataLoader{T, AT, OT, K, typeof(output_dim)}(input, output, input_dim, + input_time_steps, n_params, output_dim, output_time_steps) +end + +# The two decisions `autoencoder` makes at every constructor that takes a data set: the kind it +# builds, and the two lengths of a matrix or a `(q, p)` pair read as a set of samples, whose second +# axis is then the parameter index and whose time axis has length one. +data_set_kind(autoencoder::Bool) = autoencoder ? :RegularData : :TimeSeries +function data_set_shape(autoencoder::Bool, second_dim::Int, third_dim::Int) + autoencoder ? (1, second_dim) : (second_dim, third_dim) end """ @@ -90,7 +100,7 @@ DataLoader(data) ┌ Info: You have provided a tensor with three axes as input. They will be interpreted as └ (i) system dimension, (ii) number of time steps and (iii) number of params. -DataLoader{Int64, Array{Int64, 3}, Nothing, :TimeSeries}([1 4; 2 5; 3 6;;; 7 10; 8 11; 9 12], nothing, 3, 2, 2, nothing, nothing) +DataLoader{Int64, Array{Int64, 3}, Nothing, :TimeSeries, Nothing}([1 4; 2 5; 3 6;;; 7 10; 8 11; 9 12], nothing, 3, 2, 2, nothing, nothing) ``` But if we write @@ -107,24 +117,18 @@ DataLoader(data; suppress_info = true) # output -DataLoader{Int64, Array{Int64, 3}, Nothing, :TimeSeries}([1 4; 2 5; 3 6;;; 7 10; 8 11; 9 12], nothing, 3, 2, 2, nothing, nothing) +DataLoader{Int64, Array{Int64, 3}, Nothing, :TimeSeries, Nothing}([1 4; 2 5; 3 6;;; 7 10; 8 11; 9 12], nothing, 3, 2, 2, nothing, nothing) ``` the `@info` statement is not printed. """ function DataLoader( data::AbstractArray{T, 3}; autoencoder::Bool = false, suppress_info = false) where {T} - if !suppress_info + !suppress_info && @info "You have provided a tensor with three axes as input. They will be interpreted as \n (i) system dimension, (ii) number of time steps and (iii) number of params." - end input_dim, input_time_steps, n_params = size(data) - if autoencoder - DataLoader{T, typeof(data), Nothing, :RegularData}( - data, nothing, input_dim, input_time_steps, n_params, nothing, nothing) - else - DataLoader{T, typeof(data), Nothing, :TimeSeries}( - data, nothing, input_dim, input_time_steps, n_params, nothing, nothing) - end + DataLoader{T, typeof(data), Nothing, data_set_kind(autoencoder)}( + data, nothing, input_dim, input_time_steps, n_params, nothing, nothing) end """ @@ -138,25 +142,21 @@ See [`DataLoader(::AbstractArray{<:Number, 3})`](@ref) for details. # Implementation -Internally the data are reshaped to a tensor of shape `(size(data)..., 1)` to make for a consistent representation. +The second axis of the matrix is the number of time steps, and internally the data are reshaped to +a tensor of shape `(size(data)..., 1)`: one trajectory of one parameter. With `autoencoder = true` +the matrix is a set of samples instead, so that axis is the parameter index and the time axis has +length one. """ function DataLoader( - data::AbstractMatrix{T}; autoencoder::Bool = true, suppress_info = false) where {T} - if !suppress_info - @info "You have provided a matrix as input. The axes will be interpreted as (i) system dimension and (ii) number of parameters." - end - - if autoencoder - input_dim, n_params = size(data) - reshaped_data = reshape(data, input_dim, 1, n_params) - return DataLoader{T, typeof(reshaped_data), Nothing, :RegularData}( - reshaped_data, nothing, input_dim, 1, n_params, nothing, nothing) - else - input_dim, time_steps = size(data) - reshaped_data = reshape(data, input_dim, time_steps, 1) - return DataLoader{T, typeof(reshaped_data), Nothing, :TimeSeries}( - reshaped_data, nothing, input_dim, time_steps, 1, nothing, nothing) - end + data::AbstractMatrix{T}; autoencoder::Bool = false, suppress_info = false) where {T} + !suppress_info && + @info "You have provided a matrix as input. The axes will be interpreted as (i) system dimension and (ii) $(autoencoder ? "number of parameters" : "number of time steps")." + input_dim, second_dim = size(data) + time_steps, n_params = data_set_shape(autoencoder, second_dim, 1) + reshaped_data = reshape(data, input_dim, time_steps, n_params) + + DataLoader{T, typeof(reshaped_data), Nothing, data_set_kind(autoencoder)}( + reshaped_data, nothing, input_dim, time_steps, n_params, nothing, nothing) end """ @@ -169,20 +169,33 @@ Make an instance of `DataLoader` based on a vector. If the input to `DataLoader` is a vector, it is assumed that this vector represents one-dimensional time-series data and is therefore processed as: ```julia - DataLoader(data::AbstractVector; autoencoder=true) = DataLoader(reshape(data, 1, length(data)); autoencoder = autoencoder) + DataLoader(data::AbstractVector; autoencoder=false) = DataLoader(reshape(data, 1, length(data)); autoencoder = autoencoder) ``` """ -function DataLoader(data::AbstractVector; autoencoder = true, suppress_info = false) +function DataLoader(data::AbstractVector; autoencoder::Bool = false, suppress_info = false) DataLoader(reshape(data, 1, length(data)); autoencoder = autoencoder, suppress_info = suppress_info) end -function DataLoader(input::AbstractArray{T, 3}, output::AbstractArray{T, 3}; suppress_info = false) where {T} +# A loader that holds an output is a set of samples whatever its kind: its target is the output +# itself and never the input shifted by the sequence length. The two element types are independent, +# which is what lets a real input carry an integer one-hot target — a classifier's data, and +# `accuracy`'s argument. Two element types of which neither is an integer stay a `MethodError`. +function DataLoader(input::AbstractArray{T, 3}, output::AbstractArray{T, 3}; + suppress_info = false) where {T} + _input_and_output_loader(input, output, suppress_info) +end + +function DataLoader(input::AbstractArray{T, 3}, output::AbstractArray{T1, 3}; + suppress_info = false) where {T, T1 <: Integer} + _input_and_output_loader(input, output, suppress_info) +end + +function _input_and_output_loader(input::AbstractArray{T, 3}, output::AbstractArray{T2, 3}, + suppress_info) where {T, T2} @assert size(input, 3) == size(output, 3) - if !suppress_info - @info "You have provided an input and an output." - end + !suppress_info && @info "You have provided an input and an output." - DataLoader{T, typeof(input), typeof(output), :TimeSeries}( + DataLoader{T, typeof(input), typeof(output), :RegularData}( input, output, size(input, 1), size(input, 2), size(input, 3), size(output, 1), size(output, 2)) end @@ -212,48 +225,54 @@ See e.g. [`DataLoader{::AbstractArray{<:Number, 3}}`](@ref). function DataLoader(data::NamedTuple{(:q, :p), Tuple{AT, AT}}; autoencoder::Bool = false, suppress_info = false) where {T, AT <: AbstractMatrix{T}} - if !suppress_info + !suppress_info && @info "You have provided a NamedTuple with keys q and p; the data are matrices. This is interpreted as *symplectic data*." - end + dim2, second_dim = size(data.q) + time_steps, n_params = data_set_shape(autoencoder, second_dim, 1) + reshaped_data = (q = reshape(data.q, dim2, time_steps, n_params), + p = reshape(data.p, dim2, time_steps, n_params)) - if autoencoder - dim2, n_params = size(data.q) - reshaped_data = ( - q = reshape(data.q, dim2, 1, n_params), p = reshape(data.p, dim2, 1, n_params)) - return DataLoader{T, typeof(reshaped_data), Nothing, :RegularData}( - reshaped_data, nothing, dim2 * 2, 1, n_params, nothing, nothing) - else - dim2, time_steps = size(data.q) - reshaped_data = (q = reshape(data.q, dim2, time_steps, 1), - p = reshape(data.p, dim2, time_steps, 1)) - return DataLoader{T, typeof(reshaped_data), Nothing, :TimeSeries}( - reshaped_data, nothing, dim2 * 2, time_steps, 1, nothing, nothing) - end + DataLoader{T, typeof(reshaped_data), Nothing, data_set_kind(autoencoder)}( + reshaped_data, nothing, dim2 * 2, time_steps, n_params, nothing, nothing) end function DataLoader(data::NamedTuple{(:q, :p), Tuple{AT, AT}}; autoencoder::Bool = false, suppress_info = false) where {T, AT <: AbstractArray{T, 3}} - if !suppress_info + !suppress_info && @info "You have provided a NamedTuple with keys q and p; the data are tensors. This is interpreted as *symplectic data*." - end - dim2, time_steps, n_params = size(data.q) - if autoencoder - DataLoader{T, typeof(data), Nothing, :RegularData}( - data, nothing, dim2 * 2, time_steps, n_params, nothing, nothing) - else - DataLoader{T, typeof(data), Nothing, :TimeSeries}( - data, nothing, dim2 * 2, time_steps, n_params, nothing, nothing) - end + DataLoader{T, typeof(data), Nothing, data_set_kind(autoencoder)}( + data, nothing, dim2 * 2, time_steps, n_params, nothing, nothing) end function DataLoader(data::NamedTuple{(:q, :p), Tuple{VT, VT}}; - suppress_info = false) where {VT <: AbstractVector} + autoencoder::Bool = false, suppress_info = false) where {VT <: AbstractVector} DataLoader( (q = reshape(data.q, 1, length(data.q)), p = reshape(data.p, 1, length(data.p))); - suppress_info = suppress_info) + autoencoder = autoencoder, suppress_info = suppress_info) +end + +# One 1-based data tensor out of the coordinate of one solution. The `q` of a solution carries the +# `DataSeries` time axis `0:n` and a `DataLoader` tensor is 1-based on all three axes, so the states +# are copied in one at a time. +function data_tensor(q::AbstractVector) + data = Array{eltype(first(q))}(undef, length(first(q)), length(q), 1) + for (t, state) in enumerate(q) + data[:, t, 1] = state + end + data +end + +# The tensor of one coordinate over the members of an ensemble, along the parameter axis. +function ensemble_data_tensor(solutions, coordinate::Symbol) + first_data = data_tensor(getproperty(first(solutions), coordinate)) + data = similar(first_data, size(first_data, 1), size(first_data, 2), length(solutions)) + for (i, solution) in pairs(solutions) + data[:, :, i] = data_tensor(getproperty(solution, coordinate)) + end + data end function data_tensors_from_geometric_solution(solution::GeometricSolution{T, @@ -261,16 +280,7 @@ function data_tensors_from_geometric_solution(solution::GeometricSolution{T, TT, NT}) where {T <: Number, T2 <: Number, TT <: TimeSeries{T2}, TuT, NT <: NamedTuple{(:t, :q, :p, :q̇, :ṗ), TuT}} - sys_dim, input_time_steps = length(solution.dataser.q[0]), length(solution.t) - data = (q = zeros(T, sys_dim, input_time_steps, 1), - p = zeros(T, sys_dim, input_time_steps, 1)) - - for dim in 1:sys_dim - data.q[dim, :, 1] = solution.q[:, dim] - data.p[dim, :, 1] = solution.p[:, dim] - end - - data + (q = data_tensor(solution.q), p = data_tensor(solution.p)) end function data_tensors_from_geometric_solution(solution::GeometricSolution{T, @@ -278,14 +288,7 @@ function data_tensors_from_geometric_solution(solution::GeometricSolution{T, TT, NT}) where {T <: Number, T2 <: Number, TT <: TimeSeries{T2}, TuT, NT <: Union{NamedTuple{(:t, :q, :v), TuT}, NamedTuple{(:t, :q, :q̇), TuT}}} - sys_dim, input_time_steps = length(solution.dataser.q[0]), length(solution.t) - data = zeros(T, sys_dim, input_time_steps, 1) - - for dim in 1:sys_dim - data[dim, :, 1] = solution.q[:, dim] - end - - data + data_tensor(solution.q) end """ @@ -306,45 +309,22 @@ See the docstring for [`DataLoader(::AbstractArray{<:Number, 3})`](@ref). Internally this stores the data as a tensor where the third axis has length 1. """ -function DataLoader(solution::GeometricSolution{T, T2, TT, NT}, - suppress_info = false; - kwargs...) where - {T <: Number, - T2 <: Number, - TT <: TimeSeries{T2}, - TuT, - NT <: Union{ - NamedTuple{(:t, :q, :p, :q̇, :ṗ), TuT}, - NamedTuple{(:t, :q, :v), TuT}, - NamedTuple{(:t, :q, :q̇), TuT}}} - data = data_tensors_from_geometric_solution(solution) - - DataLoader(data; suppress_info = suppress_info, kwargs...) +function DataLoader(solution::GeometricSolution{T, T2, TT, NT}; + autoencoder::Bool = false, suppress_info = false) where {T <: Number, T2 <: Number, + TT <: TimeSeries{T2}, TuT, + NT <: Union{NamedTuple{(:t, :q, :p, :q̇, :ṗ), TuT}, + NamedTuple{(:t, :q, :v), TuT}, NamedTuple{(:t, :q, :q̇), TuT}}} + DataLoader(data_tensors_from_geometric_solution(solution); + autoencoder = autoencoder, suppress_info = suppress_info) end function DataLoader(ensemble_solution::EnsembleSolution{T, T1, Vector{ST}}; - suppress_info = false) where {T, - T1, - TuT, - TT <: TimeSeries{T1}, - ST <: Union{ + autoencoder::Bool = false, suppress_info = false) where {T, T1, TuT, + TT <: TimeSeries{T1}, ST <: Union{ GeometricSolution{T, T1, TT, NamedTuple{(:t, :q, :v), TuT}}, - GeometricSolution{T, T1, TT, NamedTuple{(:t, :q, :q̇), TuT}} - } -} - sys_dim = length(ensemble_solution.s[1].q[0]) - input_time_steps = length(ensemble_solution.t) - n_params = length(ensemble_solution.s) - - data = zeros(T, sys_dim, input_time_steps, n_params) - - for (solution, i) in zip(ensemble_solution.s, axes(ensemble_solution.s, 1)) - for dim in 1:sys_dim - data[dim, :, i] = solution.q[:, dim] - end - end - - DataLoader(data; suppress_info = suppress_info) + GeometricSolution{T, T1, TT, NamedTuple{(:t, :q, :q̇), TuT}}}} + DataLoader(ensemble_data_tensor(ensemble_solution.s, :q); + autoencoder = autoencoder, suppress_info = suppress_info) end """ @@ -367,52 +347,31 @@ See the docstring for [`DataLoader(::AbstractArray{<:Number, 3})`](@ref). Internally this stores the data as a tensor where the third axis has length equal to the number of solutions in the ensemble. """ function DataLoader(ensemble_solution::EnsembleSolution{T, T1, Vector{ST}}; - autoencoder = false, - suppress_info = false -) where {T, - T1, - TuT, - TT <: TimeSeries{T1}, - ST <: GeometricSolution{T, T1, TT, NamedTuple{(:t, :q, :p, :q̇, :ṗ), TuT}} -} - sys_dim = length(ensemble_solution.s[1].q[0]) - input_time_steps = length(ensemble_solution.t) - n_params = length(ensemble_solution.s) - - data = (q = zeros(T, sys_dim, input_time_steps, n_params), - p = zeros(T, sys_dim, input_time_steps, n_params)) - - for (solution, i) in zip(ensemble_solution.s, axes(ensemble_solution.s, 1)) - for dim in 1:sys_dim - data.q[dim, :, i] = solution.q[:, dim] - data.p[dim, :, i] = solution.p[:, dim] - end - end - - DataLoader(data; autoencoder = autoencoder, suppress_info = suppress_info) + autoencoder::Bool = false, suppress_info = false) where {T, T1, TuT, + TT <: TimeSeries{T1}, ST <: GeometricSolution{T, T1, TT, + NamedTuple{(:t, :q, :p, :q̇, :ṗ), TuT}}} + DataLoader( + (q = ensemble_data_tensor(ensemble_solution.s, :q), + p = ensemble_data_tensor(ensemble_solution.s, :p)); + autoencoder = autoencoder, suppress_info = suppress_info) end -function map_to_new_backend(input::AbstractArray{T}, backend::KernelAbstractions.Backend) where {T} - input₂ = KernelAbstractions.allocate(backend, T, size(input)...) - KernelAbstractions.copyto!(backend, input₂, input) - input₂ +function map_to_new_backend(input::AbstractArray{T}, + backend::KernelAbstractions.Backend) where {T} + converted = KernelAbstractions.allocate(backend, T, size(input)...) + KernelAbstractions.copyto!(backend, converted, input) + converted end -function map_to_new_backend(input::QPT{T}, backend::KernelAbstractions.Backend) where {T} - input₂ = (q = KernelAbstractions.allocate(backend, T, size(input.q)...), - p = KernelAbstractions.allocate(backend, T, size(input.p)...)) - KernelAbstractions.copyto!(backend, input₂.q, input.q) - KernelAbstractions.copyto!(backend, input₂.p, input.p) - input₂ +# `map` over a `(q, p)` named tuple reaches each of the two tensors; over a plain array it would +# map the array's elements instead, which is why an array has a method of its own. +function map_to_new_backend(input::QPT, backend::KernelAbstractions.Backend) + map(tensor -> map_to_new_backend(tensor, backend), input) end -function map_to_type(input::QPT, T::DataType) - (q = T.(input.q), p = T.(input.p)) -end +map_to_type(input::AbstractArray, ::Type{T}) where {T} = T.(input) -function map_to_type(input::AbstractArray, T::DataType) - T.(input) -end +map_to_type(input::QPT, ::Type{T}) where {T} = map(tensor -> map_to_type(tensor, T), input) @doc raw""" DataLoader(dl, backend) @@ -438,39 +397,17 @@ By default this inherits the autoencoder property form `dl`. See the docstring for [`DataLoader(data::AbstractArray{<:Number, 3})`](@ref). """ -function DataLoader(dl::DataLoader{T1, <:QPTOAT, Nothing, Type}, +function DataLoader(dl::DataLoader{T1, <:QPTOAT, Nothing, K}, backend::KernelAbstractions.Backend = networkbackend(dl), T::DataType = T1; autoencoder::Union{Nothing, Bool} = nothing -) where {T1, Type} - DT = if isnothing(autoencoder) - Type - elseif autoencoder - :RegularData - else - :TimeSeries - end - - input = if T == T1 - dl.input - else - map_to_type(dl.input, T) - end - - new_input = if backend == networkbackend(dl) - input - else - map_to_new_backend(input, backend) - end +) where {T1, K} + kind = isnothing(autoencoder) ? K : data_set_kind(autoencoder) + input = T == T1 ? dl.input : map_to_type(dl.input, T) + input = backend == networkbackend(dl) ? input : map_to_new_backend(input, backend) - DataLoader{T, typeof(new_input), Nothing, DT}( - new_input, - nothing, - dl.input_dim, - dl.input_time_steps, - dl.n_params, - nothing, - nothing) + DataLoader{T, typeof(input), Nothing, kind}( + input, nothing, dl.input_dim, dl.input_time_steps, dl.n_params, nothing, nothing) end function DataLoader(dl::DataLoader, T::DataType; kwargs...) diff --git a/src/data_loader/tensor_assign.jl b/src/data_loader/tensor_assign.jl index f351bc9ec..fe76ff5dc 100644 --- a/src/data_loader/tensor_assign.jl +++ b/src/data_loader/tensor_assign.jl @@ -1,21 +1,6 @@ -""" -A few data loader/tensor assignment functions. - -They assign batched tensors during training based on a tensor with three axes: (i) system dimension, (ii) parameter dependence (including initial conditions), (iii) time step. - -All of this uses KernelAbstractions and should work with any GPU supported by Julia. -""" - # Assigning a batch and its target out of the full data tensor is `src/data_loader/batch.jl`'s job: -# `assign_input_from_vector_of_tuples_kernel!` and `assign_output_from_vector_of_tuples_kernel!` -# index by a vector of `(time_step, parameter)` tuples. What is here estimates the output instead. - -@kernel function assign_output_estimate_kernel!( - output_estimate::AbstractArray{T, 3}, full_output::AbstractArray{T, 3}, - seq_length, prediction_window) where {T} - i, j, k = @index(Global, NTuple) - output_estimate[i, j, k] = full_output[i, seq_length - prediction_window + j, k] -end +# its gather kernel indexes by a vector of `(time_step, parameter)` tuples. What is here estimates +# the output instead. @doc raw""" assign_output_estimate(full_output, prediction_window) @@ -31,52 +16,17 @@ i.e. \begin{bmatrix} z^{(1)}_1 & \cdots & z^{(T)}_1 \\ \cdots & \cdots & \cdots \\ - z^{(1)}_n & \cdots & z^{(T})_n + z^{(1)}_n & \cdots & z^{(T)}_n \end{bmatrix} \mapsto \begin{bmatrix} - z^{(T - \mathtt{pw})}_1 & \cdots & z^{(T)}_1 \\ + z^{(T - \mathtt{pw} + 1)}_1 & \cdots & z^{(T)}_1 \\ \cdots & \cdots & \cdots \\ - z^{(T - \mathtt{pw})}_n & \cdots & z^{(T})_n\end{bmatrix} + z^{(T - \mathtt{pw} + 1)}_n & \cdots & z^{(T)}_n\end{bmatrix} ``` If `prediction_window` is equal to `sequence_length`, then this is not needed. """ -function assign_output_estimate(full_output::AbstractArray{T, 3}, prediction_window::Int) where {T} - sys_dim, seq_length, batch_size = size(full_output) - backend = networkbackend(full_output) - output_estimate = KernelAbstractions.allocate( - backend, T, sys_dim, prediction_window, batch_size) - assign_output_estimate! = assign_output_estimate_kernel!(networkbackend(full_output)) - assign_output_estimate!(output_estimate, full_output, seq_length, - prediction_window, ndrange = size(output_estimate)) - output_estimate -end - -# Used for differentiating `assign_output_estimate`, which appears in the loss. -@kernel function augment_zeros_kernel!( - zero_tensor::AbstractArray{T, 3}, output_diff::AbstractArray{T, 3}, - seq_length, prediction_window) where {T} - i, j, k = @index(Global, NTuple) - zero_tensor[i, seq_length - prediction_window + j, k] = output_diff[i, j, k] -end -function augment_zeros(output_diff::AbstractArray{T, 3}, seq_length) where {T} - sys_dim, prediction_window, batch_size = size(output_diff) - backend = networkbackend(output_diff) - zero_tensor = KernelAbstractions.zeros(backend, T, sys_dim, seq_length, batch_size) - augment_zeros! = augment_zeros_kernel!(networkbackend(output_diff)) - augment_zeros!(zero_tensor, output_diff, seq_length, - prediction_window, ndrange = size(output_diff)) - zero_tensor -end - -function ChainRulesCore.rrule(::typeof(assign_output_estimate), - full_output::AbstractArray{T, 3}, prediction_window) where {T} - seq_length = size(full_output, 2) - output_estimate = assign_output_estimate(full_output, prediction_window) - function assign_output_estimate_pullback(output_diff) - f̄ = NoTangent() - batch_diff = @thunk augment_zeros(unthunk(output_diff), seq_length) - return f̄, batch_diff, NoTangent() - end - return output_estimate, assign_output_estimate_pullback +function assign_output_estimate( + full_output::AbstractArray{T, 3}, prediction_window::Int) where {T} + full_output[:, (end - prediction_window + 1):end, :] end diff --git a/test/data_loader/accuracy.jl b/test/data_loader/accuracy.jl index 273fedb75..9ae757417 100644 --- a/test/data_loader/accuracy.jl +++ b/test/data_loader/accuracy.jl @@ -1,18 +1,16 @@ -using GeometricMachineLearning -using Random -using Test - # `accuracy` is documented API -- `docs/src/data_loader/data_loader.md` renders both its methods -- # and nothing in the package, the tests or the scripts called it. A documented function with no # caller and no test is a function nobody notices breaking, so this file exercises it rather than # the docs entry being dropped. # -# It also records what the audit found on the way: **no two-argument `DataLoader` constructor -# produces the type this function dispatches on.** The signature is -# `DataLoader{T, AT <: AbstractArray{T}, BT <: AbstractArray{T1}}` with `T1 <: Integer`, and every -# `DataLoader(input, output)` method requires both arrays to share one element type. A classifier's -# data are exactly the case where they differ: real inputs, integer one-hot targets. So the parametric -# constructor below is not a shortcut past a public one -- it is the only way in. +# The loader comes from `DataLoader(input, output)`, which is C17's constructor: its two element +# types are free, so a real input and an integer one-hot target are one call. Until M2 the only way +# to build this argument was to name the parametric type, and the file said so. + +using GeometricMachineLearning +using Random +using Test +using Zygote Random.seed!(1234) @@ -29,9 +27,7 @@ function classification_loader(classes::AbstractVector{<:Integer}) for (k, c) in pairs(classes) output[c, 1, k] = 1 end - GeometricMachineLearning.DataLoader{ - Float32, Array{Float32, 3}, Array{Int64, 3}, :TimeSeries}( - input, output, input_dim, seq_length, n_params, n_classes, 1) + DataLoader(input, output; suppress_info = true) end model = Chain(Dense(input_dim, n_classes, tanh)) @@ -58,3 +54,23 @@ end GeometricMachineLearning.accuracy(model, ps, dl) end end + +# Edge 10: the crop of `assign_output_estimate` is the last `prediction_window` columns, and its +# gradient is Zygote's `getindex` adjoint — zero on the cropped columns, and a cotangent of the +# primal's element type whatever the seed's. +@testset "assign_output_estimate crops the last columns ($T)" for T in (Float32, Float64) + x = T.(reshape(1:(3 * 5 * 4), 3, 5, 4)) + @test GeometricMachineLearning.assign_output_estimate(x, 2) == x[:, 4:5, :] + @test GeometricMachineLearning.assign_output_estimate(x, 5) == x + @test eltype(GeometricMachineLearning.assign_output_estimate(x, 2)) === T +end + +@testset "the gradient of assign_output_estimate is Zygote's ($T)" for T in (Float32, Float64) + x = T.(reshape(1:(3 * 5 * 4), 3, 5, 4)) + crop(z) = GeometricMachineLearning.assign_output_estimate(z, 2) + gradient = only(Zygote.gradient(z -> sum(abs2, crop(z)), x)) + @test all(iszero, gradient[:, 1:3, :]) + @test gradient[:, 4:5, :] == 2 .* x[:, 4:5, :] + _, back = Zygote.pullback(crop, x) + @test eltype(only(back(ones(Float64, 3, 2, 4)))) === T +end diff --git a/test/data_loader/autoencoder_keyword.jl b/test/data_loader/autoencoder_keyword.jl index 9081e6ac6..408dd2a70 100644 --- a/test/data_loader/autoencoder_keyword.jl +++ b/test/data_loader/autoencoder_keyword.jl @@ -25,6 +25,9 @@ end @testset "DataLoader(::NamedTuple{(:q, :p)}) of matrices" begin data = (q = rand(Float32, 2, 3), p = rand(Float32, 2, 3)) @test_throws TypeError DataLoader(data; autoencoder = nothing, suppress_info = true) + # the vector pair reaches the same branch through the matrices above + @test_throws TypeError DataLoader((q = rand(Float32, 3), p = rand(Float32, 3)); + autoencoder = nothing, suppress_info = true) end @testset "DataLoader(::NamedTuple{(:q, :p)}) of tensors" begin @@ -41,7 +44,7 @@ end dl = DataLoader(rand(Float32, 2, 3, 4); autoencoder = true, suppress_info = true) inherited = DataLoader(dl, GeometricMachineLearning.CPU()) - @test typeof(inherited).parameters[end] == typeof(dl).parameters[end] + @test typeof(inherited).parameters[4] === typeof(dl).parameters[4] === :RegularData @test_throws TypeError DataLoader(dl, GeometricMachineLearning.CPU(); autoencoder = "x") end diff --git a/test/data_loader/batch.jl b/test/data_loader/batch.jl new file mode 100644 index 000000000..9f9cd0736 --- /dev/null +++ b/test/data_loader/batch.jl @@ -0,0 +1,169 @@ +# The batching rules of `src/data_loader/batch.jl`. The data are the part's own: ten samples +# `x = T.(reshape(1:20, 2, 10))` with `y = 3 .* x`, so that `output == 3 .* input` identifies the +# pair of one sample and not the pair of another, and a drawn index pair names the sample it draws. + +using GeometricMachineLearning +using GeometricMachineLearning: convert_input_and_batch_indices_to_array, number_of_batches +using Test +import Random + +Random.seed!(1234) + +# The kind of a loader is its fourth type parameter. +kind(dl::DataLoader) = typeof(dl).parameters[4] + +samples(T) = T.(reshape(1:20, 2, 10)) + +# What one minibatch holds, for every minibatch of an epoch. +gathered(dl, batch) = map(batch(dl)) do indices + convert_input_and_batch_indices_to_array(dl, batch, indices) +end + +@testset "a supervised loader samples every pair, the last included ($T)" for T in (Float32, Float64) + x = samples(T) + y = 3 .* x + loaders = ("input and output" => DataLoader(x, y; suppress_info = true), + "a view output" => DataLoader(x, view(y, :, :); suppress_info = true)) + for (name, dl) in loaders + @test kind(dl) === :RegularData + @test (dl.input_dim, dl.input_time_steps, dl.n_params) == (2, 10, 1) + @test (dl.output_dim, dl.output_time_steps) == (2, 10) + + batch = Batch(3) + minibatches = batch(dl) + # every sample is drawn, once, the tenth included + @test sort(vcat(minibatches...)) == [(t, 1) for t in 1:10] + @test length(minibatches) == 4 + + data = gathered(dl, batch) + inputs = first.(data) + @test sort(size.(inputs, 3)) == [1, 3, 3, 3] + @test all(size(input)[1:2] == (2, 1) for input in inputs) + for (input, output) in data + @test eltype(input) === T + @test size(output) == size(input) + @test output == 3 .* input + end + # the gathered inputs are the columns of `x`, in some order + @test sort(eachcol(hcat((input[:, 1, k] + for input in inputs for k in axes(input, 3))...))) == + sort(eachcol(Matrix(x))) + end +end + +# The `(q, p)` loader of the same ten samples: `autoencoder = true` reads the second axis of a +# matrix pair as the parameter index, so its one time step carries all ten samples. +@testset "a (q, p) loader of samples draws every pair ($T)" for T in (Float32, Float64) + x = samples(T) + dl = DataLoader((q = x, p = 3 .* x); autoencoder = true, suppress_info = true) + @test kind(dl) === :RegularData + @test (dl.input_dim, dl.input_time_steps, dl.n_params) == (4, 1, 10) + + batch = Batch(3) + @test sort(vcat(batch(dl)...)) == [(1, p) for p in 1:10] + minibatches = gathered(dl, batch) + @test length(minibatches) == 4 + @test sort([size(mb.q, 3) for mb in minibatches]) == [1, 3, 3, 3] + for mb in minibatches + @test mb isa NamedTuple{(:q, :p)} + @test eltype(mb.q) === T + @test mb.p == 3 .* mb.q + end +end + +# K8: one sample is a loader the training loop can draw from. +@testset "one sample is one minibatch" begin + x = Float32.(reshape(1:2, 2, 1)) + dl = DataLoader(x, x; suppress_info = true) + batch = Batch(1) + minibatches = batch(dl) + @test length(minibatches) == 1 + @test only(minibatches) == [(1, 1)] + input, output = only(gathered(dl, batch)) + @test size(input) == (2, 1, 1) + @test output == input +end + +# Edge 3: one count, `cld(n · n_params, batch_size)` with `n` the number of start positions, and an +# `ArgumentError` where no window fits — instead of the assertion and the empty epoch before it. +@testset "one count, and a window that fits no start position" begin + trajectory = DataLoader(rand(Float64, 2, 10); suppress_info = true) + @test kind(trajectory) === :TimeSeries + @test length(Batch(3, 9, 1)(trajectory)) == 1 + @test_throws ArgumentError Batch(3, 10, 1)(trajectory) + + x = samples(Float64) + dl = DataLoader(x, 3 .* x; suppress_info = true) + @test length(Batch(3, 10, 10)(dl)) == 1 + @test_throws ArgumentError Batch(3, 11, 1)(dl) + + # the message names the time steps, the sequence length and the prediction window: three + # different numbers here, so that each one's presence is a separate reading + raised = try + Batch(3, 11, 5)(trajectory) + catch e + e + end + @test raised isa ArgumentError + for word in ("time steps", "sequence length", "prediction window") + @test occursin(word, raised.msg) + end + @test occursin("$(trajectory.input_time_steps)", raised.msg) + @test occursin("11", raised.msg) + @test occursin("5", raised.msg) +end + +# Edges 4 and 5: the classifier's loader, and `Batch(batch_size, dl)` — one draw per sample with the +# whole sequence and the whole output, which is what GMLDatasets calls. +@testset "Batch(batch_size, dl) draws every sample with the whole sequence ($T)" for T in (Float32, Float64) + n_classes, n_patches, n_images = 10, 4, 10 + input = rand(T, 9, n_patches, n_images) + output = zeros(Int, n_classes, 1, n_images) + for (k, c) in pairs(rand(1:n_classes, n_images)) + output[c, 1, k] = 1 + end + dl = DataLoader(input, output; suppress_info = true) + @test kind(dl) === :RegularData + @test (dl.input_dim, dl.input_time_steps, dl.n_params) == (9, 4, 10) + @test (dl.output_dim, dl.output_time_steps) == (10, 1) + + batch = Batch(2, dl) + @test batch == Batch(2, 4, 1) + minibatches = batch(dl) + @test length(minibatches) == 5 + @test length(Batch(2)(dl)) == 5 + for indices in minibatches + input_batch, output_batch = convert_input_and_batch_indices_to_array(dl, batch, indices) + @test size(input_batch) == (9, 4, 2) + @test size(output_batch) == (10, 1, 2) + @test eltype(input_batch) === T + @test eltype(output_batch) === Int + @test output_batch == output[:, :, [i[2] for i in indices]] + end +end + +# Edge 2: the target is the output when there is one, whatever kind the loader is typed as. A +# `DataLoader{T, A3, Array{Int, 3}, :TimeSeries}` is what GMLDatasets and two tests build by naming +# the type, and it must batch like the `:RegularData` loader of the same data. +@testset "a loader with an output is samples whatever its kind" begin + input = rand(Float32, 9, 4, 10) + output = zeros(Int, 10, 1, 10) + for (k, c) in pairs(rand(1:10, 10)) + output[c, 1, k] = 1 + end + dl = DataLoader(input, output; suppress_info = true) + dl_time_series = DataLoader{Float32, typeof(input), typeof(output), :TimeSeries}( + input, output, 9, 4, 10, 10, 1) + + for batch in (Batch(2, dl), Batch(3, 1, 1)) + Random.seed!(11) + asked = batch(dl_time_series) + Random.seed!(11) + reference = batch(dl) + @test asked == reference + for (i, j) in zip(asked, reference) + @test convert_input_and_batch_indices_to_array(dl_time_series, batch, i) == + convert_input_and_batch_indices_to_array(dl, batch, j) + end + end +end diff --git a/test/data_loader/batch_data_loader_qp_test.jl b/test/data_loader/batch_data_loader_qp_test.jl index 83fa6436e..7c2546626 100644 --- a/test/data_loader/batch_data_loader_qp_test.jl +++ b/test/data_loader/batch_data_loader_qp_test.jl @@ -21,9 +21,8 @@ function test_data_loader(dim = 2, number_of_time_steps = 100, dl2 = DataLoader(dummy_qp_data_tensor(dim, number_of_time_steps, number_of_parameters, T)) # A `(q, p)` pair is symplectic data, so the second axis is time and a pair of matrices is one - # trajectory. A bare matrix reads its second axis as the parameter index instead, which is what - # makes this path worth asserting: the same array shape means different things under the two - # constructors. + # trajectory. With `autoencoder = true` the same shape reads that axis as the parameter index + # instead, which is what makes this path worth asserting. @test size(dl1.input.q) == (dim ÷ 2, number_of_time_steps, 1) @test (dl1.input_dim, dl1.input_time_steps, dl1.n_params) == (dim, number_of_time_steps, 1) diff --git a/test/data_loader/batch_functor.jl b/test/data_loader/batch_functor.jl index 537024b84..704e9301e 100644 --- a/test/data_loader/batch_functor.jl +++ b/test/data_loader/batch_functor.jl @@ -7,7 +7,7 @@ Random.seed!(123) const data_set_raw = 1:10 const data_set = reshape(data_set_raw, 1, 10) -const dl = DataLoader(data_set) +const dl = DataLoader(data_set; autoencoder = true) function test_batching_over_one_axis(batch_size::Int = 3) batch = Batch(batch_size) diff --git a/test/data_loader/batch_index_set.jl b/test/data_loader/batch_index_set.jl index 1fcf84c83..24335f388 100644 --- a/test/data_loader/batch_index_set.jl +++ b/test/data_loader/batch_index_set.jl @@ -12,7 +12,7 @@ import Random Random.seed!(123) @test isconcretetype(Base.infer_return_type( - batch_over_two_axes, Tuple{Batch, Int, Int, DataLoader})) + batch_over_two_axes, Tuple{Batch, Int, Int})) function index_set_allocations(n_columns::Int, batch_size::Int = 8) dl = DataLoader(rand(Float32, 4, n_columns); autoencoder = false, suppress_info = true) diff --git a/test/data_loader/data_loader.jl b/test/data_loader/data_loader.jl new file mode 100644 index 000000000..f86c79b10 --- /dev/null +++ b/test/data_loader/data_loader.jl @@ -0,0 +1,162 @@ +# The constructors of `src/data_loader/data_loader.jl`: the `autoencoder` default and the keyword +# set of every method (the part's check 4), the fifth type parameter of the struct (#211), the copy +# loops out of a solution and an ensemble (B10, edge 11), the element-type conversion of +# `DataLoader(dl, T)`, and the `(input, output)` method that lets a classifier's data be built +# without naming the parametric type (C17). + +using GeometricMachineLearning +using Test +import Random + +using GeometricIntegrators: integrate, ImplicitMidpoint +using GeometricProblems.HarmonicOscillator: hodeensemble, hodeproblem +using GeometricProblems.RigidBody: odeensemble, odeproblem, default_parameters + +Random.seed!(1234) + +# The kind of a loader is its fourth type parameter. +kind(dl::DataLoader) = typeof(dl).parameters[4] + +@testset "every constructor defaults autoencoder = false ($T)" for T in (Float32, Float64) + matrix = rand(T, 3, 10) + tensor = rand(T, 3, 10, 2) + @test kind(DataLoader(matrix; suppress_info = true)) === :TimeSeries + @test kind(DataLoader(rand(T, 10); suppress_info = true)) === :TimeSeries + @test kind(DataLoader(tensor; suppress_info = true)) === :TimeSeries + @test kind(DataLoader((q = matrix, p = matrix); suppress_info = true)) === :TimeSeries + @test kind(DataLoader((q = tensor, p = tensor); suppress_info = true)) === :TimeSeries + @test kind(DataLoader((q = rand(T, 10), p = rand(T, 10)); suppress_info = true)) === + :TimeSeries + # ... and `autoencoder = true` is what makes a matrix a set of samples + @test kind(DataLoader(matrix; autoencoder = true, suppress_info = true)) === + :RegularData + @test kind(DataLoader((q = rand(T, 10), p = rand(T, 10)); + autoencoder = true, suppress_info = true)) === :RegularData + # a loader with an output is a set of samples, whatever the two element types are + @test kind(DataLoader(matrix, matrix; suppress_info = true)) === :RegularData +end + +# The matrix message names the second axis by the kind actually built. +@testset "the matrix message names the axis by the kind" begin + @test_logs (:info, r"time steps") DataLoader(rand(Float32, 3, 10)) + @test_logs (:info, r"number of parameters") DataLoader( + rand(Float32, 3, 10); autoencoder = true) +end + +# `suppress_info` is a keyword in every method of a data set, a solution and an ensemble. +@testset "suppress_info is a keyword everywhere" begin + solution = integrate( + hodeproblem(; timespan = (0.0, 1.0), timestep = 0.25), ImplicitMidpoint()) + @test DataLoader(solution; suppress_info = true) isa DataLoader + @test_throws MethodError DataLoader(solution, true) + + ensemble = integrate( + odeensemble([(q = [0.0, 0.2, 1.0],), (q = [0.1, 0.3, 1.0],)]; + timespan = (0.0, 0.4), timestep = 0.2, parameters = default_parameters()), + ImplicitMidpoint()) + @test DataLoader(ensemble; suppress_info = true) isa DataLoader + @test_throws MethodError DataLoader(ensemble, true) +end + +# B10 and edge 11: the copy loops of `data_tensors_from_geometric_solution` build one 1-based tensor +# per coordinate. The rigid-body `odeensemble` is the case whose `dataser` carries `(t, q, q̇)` and +# no `p`, which is the method B10 called unreachable. +@testset "a solution and an ensemble become 1-based tensors" begin + ics = [(q = [sin(v), 0.0, cos(v)],) for v in 0.1:0.5:2.0] + ensemble = integrate( + odeensemble(ics; timespan = (0.0, 1.0), timestep = 0.2, + parameters = default_parameters()), ImplicitMidpoint()) + dl = DataLoader(ensemble; suppress_info = true) + @test dl.input isa Array{Float64, 3} + @test kind(dl) === :TimeSeries + @test size(dl.input) == (3, length(ensemble.t), length(ensemble.s)) + for (i, solution) in pairs(ensemble.s), t in 0:(length(solution.t) - 1) + + @test dl.input[:, t + 1, i] == solution.q[t] + end + + # one solution of the same problem + solution = integrate( + odeproblem(; timespan = (0.0, 1.0), timestep = 0.2), ImplicitMidpoint()) + dl_one = DataLoader(solution; suppress_info = true) + @test dl_one.input isa Array{Float64, 3} + @test size(dl_one.input) == (3, length(solution.t), 1) + for t in 0:(length(solution.t) - 1) + @test dl_one.input[:, t + 1, 1] == solution.q[t] + end + + # and a `(q, p)` ensemble: the two coordinates are not swapped + ensemble_qp = integrate( + hodeensemble(; timespan = (0.0, 0.5), timestep = 0.25), ImplicitMidpoint()) + dl_qp = DataLoader(ensemble_qp; suppress_info = true) + @test dl_qp.input isa NamedTuple{(:q, :p)} + @test size(dl_qp.input.q) == (1, length(ensemble_qp.t), length(ensemble_qp.s)) + for (i, sol) in pairs(ensemble_qp.s), t in 0:(length(sol.t) - 1) + + @test dl_qp.input.q[:, t + 1, i] == sol.q[t] + @test dl_qp.input.p[:, t + 1, i] == sol.p[t] + end +end + +# The reproduction script `scripts/reproduction/symplectic_autoencoders/training.jl` has a `solve!` +# of a PSD architecture on a snapshot matrix, and `solve!` takes a `:RegularData` loader. That is +# why its `DataLoader(data)` passes `autoencoder = true`, and why a matrix that does not is a +# `MethodError` rather than a training run on input-shifted targets. +@testset "solve! of a PSD architecture needs a samples loader" begin + data = rand(Float32, 6, 20) + nn = NeuralNetwork(PSDArch(6, 2), CPU(), Float32) + @test isfinite(solve!(nn, DataLoader(data; autoencoder = true, suppress_info = true))) + @test_throws MethodError solve!(nn, DataLoader(data; suppress_info = true)) +end + +# #211: the two fields that are `nothing` for an unsupervised loader are concrete. +@testset "#211: the output fields are concrete in every loader" begin + loaders = (DataLoader(rand(Float32, 2, 3, 4); suppress_info = true), + DataLoader(rand(Float32, 2, 3, 4), rand(Float32, 2, 3, 4); suppress_info = true), + DataLoader(rand(Float32, 2, 3, 4); autoencoder = true, suppress_info = true), + DataLoader(rand(Float32, 2, 3); autoencoder = true, suppress_info = true)) + for dl in loaders + @test isconcretetype(fieldtype(typeof(dl), :output_dim)) + @test isconcretetype(fieldtype(typeof(dl), :output_time_steps)) + @test (fieldtype(typeof(dl), :output_dim) === Nothing) == (dl.output === nothing) + end +end + +# The conversion to another element type keeps the kind and the counts, and the other backend +# branch is the identity. +@testset "DataLoader(dl, T) converts and keeps the kind ($T)" for T in (Float32, Float64) + other = T === Float32 ? Float64 : Float32 + dl = DataLoader(rand(other, 2, 3, 4); suppress_info = true) + converted = DataLoader(dl, T) + @test eltype(converted) === T + @test converted.input isa Array{T, 3} + @test converted.input == T.(dl.input) + @test kind(converted) === kind(dl) + @test (converted.input_dim, converted.input_time_steps, converted.n_params) == + (dl.input_dim, dl.input_time_steps, dl.n_params) + @test converted.output === nothing + # the same backend and the same element type is the same tensor + @test DataLoader(dl, other).input === dl.input +end + +# C17: a classifier's real input and integer one-hot output, without naming the parametric type. +@testset "C17: a loader of an input and an integer output" begin + input = rand(Float32, 4, 3, 6) + output = zeros(Int64, 3, 1, 6) + dl = DataLoader(input, output; suppress_info = true) + @test kind(dl) === :RegularData + @test (dl.input_dim, dl.input_time_steps, dl.n_params) == (4, 3, 6) + @test (dl.output_dim, dl.output_time_steps) == (3, 1) + @test eltype(dl) === Float32 + @test dl.output === output + # the same element type in both arguments still reaches a method, and so does an `Int`/`Int` pair + @test DataLoader(input, rand(Float32, 3, 1, 6); suppress_info = true) isa DataLoader + @test DataLoader(zeros(Int, 2, 3, 4), zeros(Int64, 2, 3, 4); suppress_info = true) isa + DataLoader{Int, Array{Int, 3}, Array{Int64, 3}, :RegularData} + # two element types of which neither is an integer stay a `MethodError` + @test_throws MethodError DataLoader( + rand(Float32, 2, 3, 4), rand(Float64, 2, 3, 4); suppress_info = true) + # the third axis is still the parameter index of both + @test_throws AssertionError DataLoader( + rand(Float32, 2, 3, 4), zeros(Int, 3, 1, 5); suppress_info = true) +end diff --git a/test/data_loader/optimizer_functor_with_adam.jl b/test/data_loader/optimizer_functor_with_adam.jl index 1dd06cbe1..216a286a9 100644 --- a/test/data_loader/optimizer_functor_with_adam.jl +++ b/test/data_loader/optimizer_functor_with_adam.jl @@ -35,8 +35,10 @@ function test_optimization_with_adam(; T = Float32, input_dim = 9, n_patches = 4 T = T, input_dim = input_dim, n_patches = n_patches, n_images = n_images, n_classes = n_classes) - # batch size is equal to two - batch = Batch(2) + # The batch is drawn over the patches and the images: `Batch(2, dl)` takes the loader's own + # sequence length and output width, so one minibatch holds the whole patch sequence and the + # whole one-hot output of its two images. + batch = Batch(2, dl) # the transformer is called with `n_heads` heads and two layers model = Chain(Transformer(dl.input_dim, n_heads, 2; Stiefel = true), diff --git a/test/integration/docstrings/data_loader.jl b/test/integration/docstrings/data_loader.jl index 0526478fb..ca846430f 100644 --- a/test/integration/docstrings/data_loader.jl +++ b/test/integration/docstrings/data_loader.jl @@ -7,14 +7,16 @@ Random.seed!(1234) @testset "Data loader docstring examples" begin data = [1 2 3; 4 5 6] - @test DataLoader(data; suppress_info = true).input == reshape(data, 2, 1, 3) + @test DataLoader(data; autoencoder = true, suppress_info = true).input == + reshape(data, 2, 1, 3) + @test DataLoader(data; suppress_info = true).input == reshape(data, 2, 3, 1) tensor = [1; 2; 3;; 4; 5; 6;;; 7; 8; 9;; 10; 11; 12] expected_tensor = [1 4; 2 5; 3 6;;; 7 10; 8 11; 9 12] @test DataLoader(tensor; suppress_info = true).input == expected_tensor @test DataLoader(tensor; suppress_info = true) isa DataLoader - dl = DataLoader(rand(5); suppress_info = true) + dl = DataLoader(rand(5); autoencoder = true, suppress_info = true) batches = Batch(2)(dl) @test length(batches) == 3 @test length.(batches) == [2, 2, 1] @@ -33,10 +35,13 @@ Random.seed!(1234) @test sort(collect(union(batches₁...))) == [(i, 1) for i in 1:4] @test sort(collect(union(batches₂...))) == [(1, i) for i in 1:5] - dl = DataLoader(collect(0.1:0.1:0.9); suppress_info = true) + dl = DataLoader(collect(0.1:0.1:0.9); autoencoder = true, suppress_info = true) indices = [(1, 1), (1, 3), (1, 5)] @test convert_input_and_batch_indices_to_array(dl, Batch(3), indices) == reshape([0.1, 0.3, 0.5], 1, 1, 3) @test number_of_batches(DataLoader([1, 2, 3, 4, 5]; suppress_info = true), Batch(2)) == + 2 + @test number_of_batches( + DataLoader([1, 2, 3, 4, 5]; autoencoder = true, suppress_info = true), Batch(2)) == 3 end diff --git a/test/kernels/kernel_pullbacks.jl b/test/kernels/kernel_pullbacks.jl index a276ba318..f4815a587 100644 --- a/test/kernels/kernel_pullbacks.jl +++ b/test/kernels/kernel_pullbacks.jl @@ -1,5 +1,4 @@ using GeometricMachineLearning: tensor_tensor_mul, tensor_transpose_tensor_mul, - assign_output_estimate, tensor_mat_skew_sym_assign, cpu_inverse, tensor_inverse2, tensor_inverse3, tensor_inverse4, tensor_inverse5, init_output, @@ -23,7 +22,6 @@ function main(first_dim, second_dim, third_dim, third_tensor_dim) rand(second_dim, third_dim, third_tensor_dim)) test_rrule(tensor_transpose_tensor_mul, rand(second_dim, first_dim, third_tensor_dim), rand(second_dim, third_dim, third_tensor_dim)) - test_rrule(assign_output_estimate, rand(first_dim, second_dim, third_tensor_dim), 1) test_rrule( tensor_mat_skew_sym_assign, rand(first_dim, second_dim, third_tensor_dim + 1), rand(first_dim, first_dim)) diff --git a/test/quality/jet.jl b/test/quality/jet.jl index 98873f94b..5d0a79ec9 100644 --- a/test/quality/jet.jl +++ b/test/quality/jet.jl @@ -116,23 +116,21 @@ const A3{T} = Array{T, 3} end # src/data_loader/batch.jl - qp = (q = zeros(2, 8, 4), p = zeros(2, 8, 4)) indices = ones(Int, 2, 4) TimeSeriesQP = DataLoader{ - Float64, @NamedTuple{q::A3{Float64}, p::A3{Float64}}, Nothing, :TimeSeries} + Float64, @NamedTuple{q::A3{Float64}, p::A3{Float64}}, Nothing, :TimeSeries, Nothing} Indices = Vector{Tuple{Int, Int}} # test/data_loader/draw_batch_for_tensor_test.jl @test isempty(JET.get_reports(JET.report_opt( GML.convert_input_and_batch_indices_to_array, (TimeSeriesQP, Batch{:Transformer}, Indices); target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.assign_input_from_vector_of_tuples_kernel!, - (2, 3, 4), zeros(2, 3, 4), zeros(2, 3, 4), qp, indices)) - @test isempty(kernel_body_reports(GML.assign_output_from_vector_of_tuples_kernel!, - (2, 3, 4), zeros(2, 3, 4), zeros(2, 3, 4), qp, indices, 3)) + # the one gather kernel of a `(q, p)` input: `gather` launches it once per tensor + @test isempty(kernel_body_reports(GML.assign_from_vector_of_tuples_kernel!, + (2, 3, 4), zeros(2, 3, 4), zeros(2, 8, 4), indices, 3)) # test/integration/docstrings/data_loader.jl @test isempty(JET.get_reports(JET.report_opt( GML.convert_input_and_batch_indices_to_array, - (DataLoader{Float64, A3{Float64}, Nothing, :RegularData}, + (DataLoader{Float64, A3{Float64}, Nothing, :RegularData, Nothing}, Batch{:FeedForward}, Indices); target_modules = TARGET))) for (T, BatchType) in ((Float32, Batch{:Transformer}), ( @@ -140,21 +138,19 @@ const A3{T} = Array{T, 3} # the optimizer functor in test/data_loader/training_history_eltype.jl @test isempty(JET.get_reports(JET.report_opt( GML.convert_input_and_batch_indices_to_array, - (DataLoader{T, A3{T}, Nothing, :TimeSeries}, Batch{:FeedForward}, Indices); + (DataLoader{T, A3{T}, Nothing, :TimeSeries, Nothing}, + Batch{:FeedForward}, + Indices); target_modules = TARGET))) # input and output: test/data_loader/data_loader_for_input_and_output.jl # (`Float32`, `Batch(10, 1, 1)`) and # test/architectures/lagrangian_neural_network_tests.jl (`Float64`, `Batch(10)`) - reports = JET.get_reports(JET.report_opt( + @test isempty(JET.get_reports(JET.report_opt( GML.convert_input_and_batch_indices_to_array, - (DataLoader{T, A3{T}, A3{T}, :TimeSeries}, BatchType, Indices); - target_modules = TARGET)) - @test_broken isempty(reports) # #325 - @test isempty(kernel_body_reports( - GML.assign_input_from_vector_of_tuples_kernel!, - (2, 3, 4), zeros(T, 2, 3, 4), zeros(T, 2, 8, 4), indices)) - @test isempty(kernel_body_reports( - GML.assign_output_from_vector_of_tuples_kernel!, + (DataLoader{T, A3{T}, A3{T}, :TimeSeries, Int}, BatchType, Indices); + target_modules = TARGET))) + # the gather kernel, once per element type + @test isempty(kernel_body_reports(GML.assign_from_vector_of_tuples_kernel!, (2, 3, 4), zeros(T, 2, 3, 4), zeros(T, 2, 8, 4), indices, 3)) end # the `Batch` functor under `@allocated` in test/data_loader/batch_index_set.jl, and @@ -162,25 +158,27 @@ const A3{T} = Array{T, 3} # test/integration/docstrings/data_loader.jl, `Float64` in # test/data_loader/draw_batch_for_tensor_test.jl @test isempty(JET.get_reports(JET.report_opt(Batch(8), - (DataLoader{Float32, A3{Float32}, Nothing, :TimeSeries},); target_modules = TARGET))) + (DataLoader{Float32, A3{Float32}, Nothing, :TimeSeries, Nothing},); + target_modules = TARGET))) @test isempty(JET.get_reports(JET.report_opt(Batch(3), - (DataLoader{Int, A3{Int}, Nothing, :TimeSeries},); target_modules = TARGET))) + (DataLoader{Int, A3{Int}, Nothing, :TimeSeries, Nothing},); + target_modules = TARGET))) @test isempty(JET.get_reports(JET.report_opt(Batch(10, 16), (TimeSeriesQP,); target_modules = TARGET))) + # `Batch(batch_size, dl)` reads the loader's own two lengths, and + # test/data_loader/batch.jl calls it at the element types below + @test isempty(JET.get_reports(JET.report_opt(Batch(2), + (DataLoader{Float32, A3{Float32}, A3{Int}, :RegularData, Int},); + target_modules = TARGET))) + @test isempty(JET.get_reports(JET.report_opt(Batch(2), + (DataLoader{Float32, A3{Float32}, A3{Float32}, :TimeSeries, Int},); + target_modules = TARGET))) # src/data_loader/tensor_assign.jl: `assign_output_estimate` in - # test/kernels/kernel_pullbacks.jl (`Float64`) and test/data_loader/accuracy.jl - # (`Float32`); `augment_zeros` through the pullback of `assign_output_estimate` + # test/data_loader/accuracy.jl, in both precisions for T in (Float32, Float64) @test isempty(JET.get_reports(JET.report_opt(GML.assign_output_estimate, (A3{T}, Int); target_modules = TARGET))) - @test isempty(kernel_body_reports( - GML.assign_output_estimate_kernel!, (2, 1, 3), - zeros(T, 2, 1, 3), zeros(T, 2, 4, 3), 4, 1)) - @test isempty(JET.get_reports(JET.report_opt(GML.augment_zeros, (A3{T}, Int); - target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.augment_zeros_kernel!, (2, 1, 3), - zeros(T, 2, 4, 3), zeros(T, 2, 1, 3), 4, 1)) end else @test_skip "JET does not work on Julia $(VERSION)" # aviatesk/JET.jl#681 diff --git a/test/runtests.jl b/test/runtests.jl index 60bcb4856..ebc25815a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -83,6 +83,8 @@ if "core" in GROUPS @safetestset "train! takes a TrainingOptimizer's algorithm and linesearch" include("optimizers/optimizer.jl") @safetestset "The minibatch index set is concrete and grows linearly" include("data_loader/batch_index_set.jl") @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 "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 f364f6219a7d302dc4ad58d72e96931aa577e103 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 16 Oct 2026 08:05:48 +0200 Subject: [PATCH 2/8] CHANGELOG: the DataLoader part of 0.9 The autoencoder default, the supervised kind, the fifth type parameter and the new element-type constructor; B10, C17, K8, K12 and the batch half of #325. --- CHANGELOG.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7dc1acf26..ee24cecc5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,8 +37,14 @@ breaking release). - **The three 3-tensor functors of `Dense` and `Linear` and the HDF5 `save` and `load` methods on `NeuralNetwork` are removed**, because AbstractNeuralNetworks 0.9 provides them. The HDF5 extension and the HDF5 weak dependency are gone; with HDF5 loaded, AbstractNeuralNetworks' extension provides `save` and `load`. Its `load` has no `backend` keyword: load on the CPU, then call `changebackend(backend, nn)`. Aqua's `piracies` check runs again with 0 piracies (was 3), and the ambiguity count is 0 (was 1). This closes issue #145 and known issue B7. +- **`Batch(::Int, ::Nothing, ::Int)` is removed.** It existed only to raise `"Cannot provide prediction window alone. Need sequence length!"`, and `Batch(batch_size, sequence_length, prediction_window)` is unchanged, so a caller of that form gets a `MethodError` where it got the error. + ### Changed (breaking) +- **`DataLoader`'s `autoencoder` default is `false` in every constructor, and a loader given an input and an output is a set of samples.** A matrix or a vector is a time series unless `autoencoder = true` is passed: that keyword is what reads a matrix's second axis as the parameter index, and the matrix and the vector method defaulted to `true`, so migrate `DataLoader(matrix)` to `DataLoader(matrix; autoencoder = true)` where the columns are samples — `scripts/reproduction/symplectic_autoencoders/training.jl` and the tutorials do. `DataLoader(input, output)` builds a `:RegularData` loader whatever its kind, so its target is its output and never the input shifted by the sequence length: at the base a `DataLoader(x, y)` of ten pairs drew the samples 1 to 9 and, where the two element types differed, returned one minibatch holding the whole data set. The count is `cld(n · n_params, batch_size)` with `n` the start positions of one parameter index, and a loader and a `Batch` that leave no start position raise an `ArgumentError` naming the time steps, the sequence length and the prediction window, where the base asserted and an epoch could be empty. `Batch(batch_size, dl)` takes its two lengths from the loader, so a classifier's minibatch holds the whole patch sequence and the whole one-hot output of each of its samples. `suppress_info` is a keyword in every method, which makes `DataLoader(solution, true)` a `MethodError`. + +- **The two output fields of `DataLoader` carry a fifth type parameter.** `output_dim` and `output_time_steps` are the type `S <: Union{Int, Nothing}` of the loader, so `DataLoader{T, AT, OT, :TimeSeries}` still names a loader and a *full* type name gains `, Nothing` for a loader without an output; a four-parameter constructor `DataLoader{T, AT, OT, K}(input, output, input_dim, input_time_steps, n_params, output_dim, output_time_steps)` fills `S` in from the fields. + - **`ZygotePullback`'s `back(Δ)` is a `NetworkParameters`, not a one-tuple of one, as `SymbolicPullback`'s already was.** A structured leaf holds its storage gradient: a `SymmetricMatrix` leaf comes back as a `SymmetricMatrix` whose off-diagonal entries are twice those of the projection of the dense gradient, which is the quantity an optimizer step needs (issue #319, closed). Both `ZygotePullback` methods take `ps::NetworkParameters`, so a bare `NamedTuple` of parameters is not a `MethodError`: it misses both and reaches AbstractNeuralNetworks' untyped `AbstractPullback` fallback, which raises `ErrorException("Pullback not implemented for input-output pair!")` for an (input, output) pair and `ErrorException("Pullback not implemented for single input!")` for a single input. - **A loss functor takes a `NetworkParameters`.** The `Union{NetworkParameters, NamedTuple}` arm of `HNNLoss`, `LNNLoss`, `SymplecticEulerLoss` and `VariationalMidpointLoss` is gone, and `AutoEncoderLoss`, `TransformerLoss` and `ReducedLoss` annotate theirs `::NetworkParameters`. NeuralNetworkParameters 0.5's Zygote rule keeps the wrapper on the forward pass, so the annotation is dispatchable under AD and a gradient set of another shape can no longer reach a loss silently. Migrate a custom loss's `params::Union{NetworkParameters, NamedTuple}` to `ps::NetworkParameters`, and a hand-rolled pullback that unwraps its result to one that does not. @@ -71,6 +77,10 @@ breaking release). - **A network with a Stiefel leaf trains on data of the other precision.** Training on data whose element type differs from the network's raised `AssertionError: eltype(A) == eltype(B)` in `mat_tensor_mul`. `_mul` checks no element type, because GEMM promotes. This closes known issue K7, which leaves `KNOWN_ISSUES.md`. The test is `test/integration/mixed_precision_training.jl`. +- **A `DataLoader`'s two output fields are concrete, its minibatch index set is concrete and draws every sample, and a loader of one sample batches.** `output_dim` and `output_time_steps` were `Union{Int, Nothing}` whatever the loader held, which made a kernel launch read a value of no single type (issue #211); the copy loops of a solution and an ensemble build a 1-based tensor, and the batch indices are one host-built `2 × n` matrix copied to the backend once, with no `@views` (issue #212, which asked for their removal). A `DataLoader` of one sample drew no minibatch (known issue K8) and a loader of matrices never drew its last sample (known issue K12); both now draw every pair, and the count of a trajectory is `cld(n · n_params, batch_size)` in the two files, so a seeded run draws the minibatches it drew before, except the ones the last sample joins. The `JET.report_opt` line of `test/quality/jet.jl` for a launcher that reads a `Union{Int, Nothing}` field is a `@test` rather than a `@test_broken`: the batch half of issue #325 is fixed here. This closes known issues B10, C17, K8 and K12, which leave `KNOWN_ISSUES.md`. B10's `DataLoader(::EnsembleSolution{T, T1, Vector{ST}})` for a `dataser` of `(t, q, q̇)` is reachable — a rigid-body `odeensemble` produces one, and the tutorial and three scripts call it — and `test/data_loader/data_loader.jl` tests it, so the method stays. + +- **`assign_output_estimate` is a crop of the last `prediction_window` columns, and its gradient is Zygote's `getindex` adjoint.** Its kernel, `augment_zeros` and the `rrule` are gone: the adjoint is zero on the cropped columns, runs on a device with scalar indexing off, and gives a `Float32` cotangent for a `Float64` seed on `Float32` data, where `augment_zeros` returned the seed's type. The docstring's formula starts at column `T − pw + 1`, which is what the code does. + ### Added - **`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. @@ -84,6 +94,8 @@ breaking release). - **`Pkg.test` gains a `JLArray` device test in the `core` group.** `test/devices/jlarray.jl` runs `tensor_tensor_mul` and its pullback, `_mul` with each structured weight in the forward pass and in the storage gradient, and `assign_q_and_p` on a `JLArray` under `allowscalar(false)`. `JLArrays`, `GPUArraysCore` and `ChainRulesCore` join `test/Project.toml`, with `JLArrays = "0.3"`. `NNlib.batched_mul` is not used for `tensor_tensor_mul`, because on a backend with no NNlib extension, JLArray and Metal among them, it falls back to scalar indexing. +- **`DataLoader(input, output)` accepts two element types, one of them an integer.** `DataLoader(input::AbstractArray{T, 3}, output::AbstractArray{<:Integer, 3})` sits beside the method whose two arguments share one element type, so a classifier's data — a real input, an integer one-hot target — are one call and `accuracy` is reachable without naming the parametric type (known issue C17, which leaves `KNOWN_ISSUES.md`). The two element types are independent, and two of which neither is an integer stay a `MethodError`. + ### Changed - `Pkg.test()` no longer runs the doctests. `test/quality/doctests.jl` is now the group `doctests`, which an empty `ARGS` does not run; `Pkg.test(test_args = ["doctests"])` runs it. In CI the Doctests job stays their runner, so the test matrix no longer runs them a second time. From 616919e5b75ed32e2e068e8579769fda6c8764b0 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 16 Oct 2026 08:25:43 +0200 Subject: [PATCH 3/8] DataLoader: the ruled cuts for check 6, to 272 code lines The advisor ruling of 2026-10-16: the index set through Iterators.partition, the kernel signature without its type parameters, the index matrix in one comprehension, one if/elseif/else in the launcher, one data_set_shape, one (input, output) method with T1 <: Union{T, Integer}, the copy methods dispatching on the dataser keys, the matrix and (q, p) methods forwarding to the tensor method, if !suppress_info at the five sites, and ::Type{T} at the last two. --- CHANGELOG.md | 2 +- src/data_loader/batch.jl | 37 ++++++---------- src/data_loader/data_loader.jl | 77 ++++++++++++++-------------------- test/data_loader/batch.jl | 1 + 4 files changed, 46 insertions(+), 71 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ee24cecc5..0dd0df151 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -94,7 +94,7 @@ breaking release). - **`Pkg.test` gains a `JLArray` device test in the `core` group.** `test/devices/jlarray.jl` runs `tensor_tensor_mul` and its pullback, `_mul` with each structured weight in the forward pass and in the storage gradient, and `assign_q_and_p` on a `JLArray` under `allowscalar(false)`. `JLArrays`, `GPUArraysCore` and `ChainRulesCore` join `test/Project.toml`, with `JLArrays = "0.3"`. `NNlib.batched_mul` is not used for `tensor_tensor_mul`, because on a backend with no NNlib extension, JLArray and Metal among them, it falls back to scalar indexing. -- **`DataLoader(input, output)` accepts two element types, one of them an integer.** `DataLoader(input::AbstractArray{T, 3}, output::AbstractArray{<:Integer, 3})` sits beside the method whose two arguments share one element type, so a classifier's data — a real input, an integer one-hot target — are one call and `accuracy` is reachable without naming the parametric type (known issue C17, which leaves `KNOWN_ISSUES.md`). The two element types are independent, and two of which neither is an integer stay a `MethodError`. +- **`DataLoader(input, output)` accepts two element types, one of them an integer.** It is one method, `DataLoader(input::AbstractArray{T, 3}, output::AbstractArray{T1, 3}) where {T, T1 <: Union{T, Integer}}`, so the output's element type is the input's or an integer: a classifier's data — a real input, an integer one-hot target — are one call and `accuracy` is reachable without naming the parametric type (known issue C17, which leaves `KNOWN_ISSUES.md`). Any other pair of element types is a `MethodError`. ### Changed diff --git a/src/data_loader/batch.jl b/src/data_loader/batch.jl index a9113c511..4a27344be 100644 --- a/src/data_loader/batch.jl +++ b/src/data_loader/batch.jl @@ -184,36 +184,24 @@ end function batch_over_two_axes(batch::Batch, number_columns::Int, third_dim::Int) complete_indices = vec( collect(Iterators.product(shuffle(1:number_columns), shuffle(1:third_dim)))) - batches = Vector{Vector{Tuple{Int, Int}}}() - for i in 1:cld(length(complete_indices), batch.batch_size) - push!(batches, - complete_indices[((i - 1) * batch.batch_size + 1):min(i * batch.batch_size, end)]) - end - batches + map(collect, Iterators.partition(complete_indices, batch.batch_size)) end function (batch::Batch)(dl::DataLoader) batch_over_two_axes(batch, number_of_start_positions(dl, batch), dl.n_params) end -@kernel function assign_from_vector_of_tuples_kernel!( - input::AT, data::BT, indices::AbstractArray{Int, 2}, - offset) where { - T, AT <: AbstractArray{T, 3}, BT <: AbstractArray{T, 3}} +@kernel function assign_from_vector_of_tuples_kernel!(gathered, data, indices, offset) i, j, k = @index(Global, NTuple) - input[i, j, k] = data[i, indices[1, k] + offset + j - 1, indices[2, k]] + gathered[i, j, k] = data[i, indices[1, k] + offset + j - 1, indices[2, k]] end # A kernel cannot index a `Vector{Tuple{Int, Int}}`, so the batch indices go in as a `2 × n` matrix: # the host builds it and the backend takes it in one copy. function convert_vector_of_tuples_to_matrix(backend::Backend, batch_indices_tuple::Vector{Tuple{Int, Int}}) - batch_indices = zeros(Int, 2, length(batch_indices_tuple)) - for t in axes(batch_indices_tuple, 1) - batch_indices[1, t], batch_indices[2, t] = batch_indices_tuple[t] - end - map_to_new_backend(batch_indices, backend) + map_to_new_backend([index[i] for i in 1:2, index in batch_indices_tuple], backend) end # The gather kernel, applied to a tensor or to the `q` and `p` of a `(q, p)` pair. The gathered @@ -271,15 +259,16 @@ function convert_input_and_batch_indices_to_array( batch_size = length(batch_indices_tuple) indices = convert_vector_of_tuples_to_matrix(networkbackend(dl), batch_indices_tuple) input = gather(dl.input, dl.input_dim, indices, 0, batch.seq_length, batch_size) - # a loader that holds an output is a set of samples whatever its kind: the target is its output + # a loader that holds an output is a set of samples whatever its kind: the target is its output; + # without one, a time series predicts its own continuation and a set of samples has no target if dl.output !== nothing - return input, - gather(dl.output, dl.output_dim, indices, 0, batch.prediction_window, + input, + gather(dl.output, dl.output_dim, indices, 0, batch.prediction_window, batch_size) + elseif K === :TimeSeries + input, + gather(dl.input, dl.input_dim, indices, batch.seq_length, batch.prediction_window, batch_size) + else + input end - # without one, a time series predicts its own continuation and a set of samples has no target - K === :TimeSeries || return input - input, - gather(dl.input, dl.input_dim, indices, batch.seq_length, batch.prediction_window, - batch_size) end diff --git a/src/data_loader/data_loader.jl b/src/data_loader/data_loader.jl index cc56b417f..33b26f9cd 100644 --- a/src/data_loader/data_loader.jl +++ b/src/data_loader/data_loader.jl @@ -61,9 +61,7 @@ end # builds, and the two lengths of a matrix or a `(q, p)` pair read as a set of samples, whose second # axis is then the parameter index and whose time axis has length one. data_set_kind(autoencoder::Bool) = autoencoder ? :RegularData : :TimeSeries -function data_set_shape(autoencoder::Bool, second_dim::Int, third_dim::Int) - autoencoder ? (1, second_dim) : (second_dim, third_dim) -end +data_set_shape(autoencoder::Bool, n::Int) = autoencoder ? (1, n) : (n, 1) """ DataLoader(data::AbstractArray{<:Number, 3}) @@ -123,8 +121,9 @@ the `@info` statement is not printed. """ function DataLoader( data::AbstractArray{T, 3}; autoencoder::Bool = false, suppress_info = false) where {T} - !suppress_info && + if !suppress_info @info "You have provided a tensor with three axes as input. They will be interpreted as \n (i) system dimension, (ii) number of time steps and (iii) number of params." + end input_dim, input_time_steps, n_params = size(data) DataLoader{T, typeof(data), Nothing, data_set_kind(autoencoder)}( @@ -149,14 +148,13 @@ length one. """ function DataLoader( data::AbstractMatrix{T}; autoencoder::Bool = false, suppress_info = false) where {T} - !suppress_info && + if !suppress_info @info "You have provided a matrix as input. The axes will be interpreted as (i) system dimension and (ii) $(autoencoder ? "number of parameters" : "number of time steps")." + end input_dim, second_dim = size(data) - time_steps, n_params = data_set_shape(autoencoder, second_dim, 1) - reshaped_data = reshape(data, input_dim, time_steps, n_params) - - DataLoader{T, typeof(reshaped_data), Nothing, data_set_kind(autoencoder)}( - reshaped_data, nothing, input_dim, time_steps, n_params, nothing, nothing) + time_steps, n_params = data_set_shape(autoencoder, second_dim) + DataLoader(reshape(data, input_dim, time_steps, n_params); + autoencoder = autoencoder, suppress_info = true) end """ @@ -177,23 +175,15 @@ function DataLoader(data::AbstractVector; autoencoder::Bool = false, suppress_in end # A loader that holds an output is a set of samples whatever its kind: its target is the output -# itself and never the input shifted by the sequence length. The two element types are independent, -# which is what lets a real input carry an integer one-hot target — a classifier's data, and -# `accuracy`'s argument. Two element types of which neither is an integer stay a `MethodError`. -function DataLoader(input::AbstractArray{T, 3}, output::AbstractArray{T, 3}; - suppress_info = false) where {T} - _input_and_output_loader(input, output, suppress_info) -end - +# itself and never the input shifted by the sequence length. The output's element type is the +# input's or an integer, so a real input can carry an integer one-hot target — a classifier's data, +# and `accuracy`'s argument — and any other pair of element types is a `MethodError`. function DataLoader(input::AbstractArray{T, 3}, output::AbstractArray{T1, 3}; - suppress_info = false) where {T, T1 <: Integer} - _input_and_output_loader(input, output, suppress_info) -end - -function _input_and_output_loader(input::AbstractArray{T, 3}, output::AbstractArray{T2, 3}, - suppress_info) where {T, T2} + suppress_info = false) where {T, T1 <: Union{T, Integer}} @assert size(input, 3) == size(output, 3) - !suppress_info && @info "You have provided an input and an output." + if !suppress_info + @info "You have provided an input and an output." + end DataLoader{T, typeof(input), typeof(output), :RegularData}( input, output, size(input, 1), size(input, 2), @@ -225,22 +215,21 @@ See e.g. [`DataLoader{::AbstractArray{<:Number, 3}}`](@ref). function DataLoader(data::NamedTuple{(:q, :p), Tuple{AT, AT}}; autoencoder::Bool = false, suppress_info = false) where {T, AT <: AbstractMatrix{T}} - !suppress_info && + if !suppress_info @info "You have provided a NamedTuple with keys q and p; the data are matrices. This is interpreted as *symplectic data*." + end dim2, second_dim = size(data.q) - time_steps, n_params = data_set_shape(autoencoder, second_dim, 1) - reshaped_data = (q = reshape(data.q, dim2, time_steps, n_params), - p = reshape(data.p, dim2, time_steps, n_params)) - - DataLoader{T, typeof(reshaped_data), Nothing, data_set_kind(autoencoder)}( - reshaped_data, nothing, dim2 * 2, time_steps, n_params, nothing, nothing) + time_steps, n_params = data_set_shape(autoencoder, second_dim) + DataLoader(map(tensor -> reshape(tensor, dim2, time_steps, n_params), data); + autoencoder = autoencoder, suppress_info = true) end function DataLoader(data::NamedTuple{(:q, :p), Tuple{AT, AT}}; autoencoder::Bool = false, suppress_info = false) where {T, AT <: AbstractArray{T, 3}} - !suppress_info && + if !suppress_info @info "You have provided a NamedTuple with keys q and p; the data are tensors. This is interpreted as *symplectic data*." + end dim2, time_steps, n_params = size(data.q) DataLoader{T, typeof(data), Nothing, data_set_kind(autoencoder)}( @@ -275,19 +264,15 @@ function ensemble_data_tensor(solutions, coordinate::Symbol) data end -function data_tensors_from_geometric_solution(solution::GeometricSolution{T, - T2, - TT, - NT}) where {T <: Number, T2 <: Number, TT <: TimeSeries{T2}, - TuT, NT <: NamedTuple{(:t, :q, :p, :q̇, :ṗ), TuT}} +function data_tensors_from_geometric_solution( + solution::GeometricSolution{ + <:Number, <:Number, <:Any, <:NamedTuple{(:t, :q, :p, :q̇, :ṗ)}}) (q = data_tensor(solution.q), p = data_tensor(solution.p)) end -function data_tensors_from_geometric_solution(solution::GeometricSolution{T, - T2, - TT, - NT}) where {T <: Number, T2 <: Number, TT <: TimeSeries{T2}, TuT, - NT <: Union{NamedTuple{(:t, :q, :v), TuT}, NamedTuple{(:t, :q, :q̇), TuT}}} +function data_tensors_from_geometric_solution(solution::GeometricSolution{ + <:Number, <:Number, + <:Any, <:Union{NamedTuple{(:t, :q, :v)}, NamedTuple{(:t, :q, :q̇)}}}) data_tensor(solution.q) end @@ -399,9 +384,9 @@ See the docstring for [`DataLoader(data::AbstractArray{<:Number, 3})`](@ref). """ function DataLoader(dl::DataLoader{T1, <:QPTOAT, Nothing, K}, backend::KernelAbstractions.Backend = networkbackend(dl), - T::DataType = T1; + ::Type{T} = T1; autoencoder::Union{Nothing, Bool} = nothing -) where {T1, K} +) where {T1, K, T} kind = isnothing(autoencoder) ? K : data_set_kind(autoencoder) input = T == T1 ? dl.input : map_to_type(dl.input, T) input = backend == networkbackend(dl) ? input : map_to_new_backend(input, backend) @@ -410,7 +395,7 @@ function DataLoader(dl::DataLoader{T1, <:QPTOAT, Nothing, K}, input, nothing, dl.input_dim, dl.input_time_steps, dl.n_params, nothing, nothing) end -function DataLoader(dl::DataLoader, T::DataType; kwargs...) +function DataLoader(dl::DataLoader, ::Type{T}; kwargs...) where {T} DataLoader(dl, networkbackend(dl), T; kwargs...) end diff --git a/test/data_loader/batch.jl b/test/data_loader/batch.jl index 9f9cd0736..9979f15a7 100644 --- a/test/data_loader/batch.jl +++ b/test/data_loader/batch.jl @@ -34,6 +34,7 @@ end # every sample is drawn, once, the tenth included @test sort(vcat(minibatches...)) == [(t, 1) for t in 1:10] @test length(minibatches) == 4 + @test number_of_batches(dl, batch) == 4 data = gathered(dl, batch) inputs = first.(data) From a5e959e00ef86eab4afb1291c6de3f28d02bc5eb Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 16 Oct 2026 11:37:14 +0200 Subject: [PATCH 4/8] DataLoader: round 1's defects, the device gather and the converted counts - `map_to_new_backend` copies with `Base.copyto!`, so the host index matrix reaches a JLArray: `KernelAbstractions.copyto!` has a method for the CPU alone. - the four-parameter constructor converts its five counts, so the `Int32` sizes `GMLDatasets` passes build a loader. - `check_window` names the input's and the output's time steps and the parameter count, and it rejects an empty parameter axis, which drew an empty epoch in silence. - `ensemble_data_tensor` fills the final tensor through one coordinate copy, so no whole coordinate tensor is built per member. - the tests assert the element type of a supervised minibatch and of a `(q, p)` target, compare the gathered columns of a `(q, p)` loader and of one `(q, p)` solution with their own data, and run the one-sample and typed-kind rows in both precisions. - the comments state the current code, and the manual states the `autoencoder` default correctly. - `KNOWN_ISSUES.md` gains K18, the tensor-sampling formula that is one short. --- CHANGELOG.md | 4 +- KNOWN_ISSUES.md | 16 +++++++ docs/src/data_loader/data_loader.md | 2 +- src/data_loader/batch.jl | 23 +++++----- src/data_loader/data_loader.jl | 44 +++++++++++-------- test/data_loader/accuracy.jl | 11 +++-- test/data_loader/batch.jl | 66 +++++++++++++++++++++-------- test/data_loader/data_loader.jl | 47 ++++++++++++++++---- 8 files changed, 149 insertions(+), 64 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0dd0df151..029b48376 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -41,7 +41,7 @@ breaking release). ### Changed (breaking) -- **`DataLoader`'s `autoencoder` default is `false` in every constructor, and a loader given an input and an output is a set of samples.** A matrix or a vector is a time series unless `autoencoder = true` is passed: that keyword is what reads a matrix's second axis as the parameter index, and the matrix and the vector method defaulted to `true`, so migrate `DataLoader(matrix)` to `DataLoader(matrix; autoencoder = true)` where the columns are samples — `scripts/reproduction/symplectic_autoencoders/training.jl` and the tutorials do. `DataLoader(input, output)` builds a `:RegularData` loader whatever its kind, so its target is its output and never the input shifted by the sequence length: at the base a `DataLoader(x, y)` of ten pairs drew the samples 1 to 9 and, where the two element types differed, returned one minibatch holding the whole data set. The count is `cld(n · n_params, batch_size)` with `n` the start positions of one parameter index, and a loader and a `Batch` that leave no start position raise an `ArgumentError` naming the time steps, the sequence length and the prediction window, where the base asserted and an epoch could be empty. `Batch(batch_size, dl)` takes its two lengths from the loader, so a classifier's minibatch holds the whole patch sequence and the whole one-hot output of each of its samples. `suppress_info` is a keyword in every method, which makes `DataLoader(solution, true)` a `MethodError`. +- **`DataLoader`'s `autoencoder` default is `false` in every constructor, and a loader given an input and an output is a set of samples.** A matrix or a vector is a time series unless `autoencoder = true` is passed: that keyword is what reads a matrix's second axis as the parameter index, and the matrix and the vector method defaulted to `true`, so migrate `DataLoader(matrix)` to `DataLoader(matrix; autoencoder = true)` where the columns are samples — `scripts/reproduction/symplectic_autoencoders/training.jl` and the tutorials do. `DataLoader(input, output)` builds a `:RegularData` loader whatever its kind, so its target is its output and never the input shifted by the sequence length: at the base a `DataLoader(x, y)` of ten pairs drew the samples 1 to 9 and, where the input and the output differed in array type, returned one minibatch holding the whole data set. The count is `cld(n · n_params, batch_size)` with `n` the start positions of one parameter index, and a loader and a `Batch` that leave no start position raise an `ArgumentError` naming the time steps, the sequence length and the prediction window, where the base asserted and an epoch could be empty. `Batch(batch_size, dl)` takes its two lengths from the loader, so a classifier's minibatch holds the whole patch sequence and the whole one-hot output of each of its samples. `suppress_info` is a keyword in every method, which makes `DataLoader(solution, true)` a `MethodError`. - **The two output fields of `DataLoader` carry a fifth type parameter.** `output_dim` and `output_time_steps` are the type `S <: Union{Int, Nothing}` of the loader, so `DataLoader{T, AT, OT, :TimeSeries}` still names a loader and a *full* type name gains `, Nothing` for a loader without an output; a four-parameter constructor `DataLoader{T, AT, OT, K}(input, output, input_dim, input_time_steps, n_params, output_dim, output_time_steps)` fills `S` in from the fields. @@ -77,7 +77,7 @@ breaking release). - **A network with a Stiefel leaf trains on data of the other precision.** Training on data whose element type differs from the network's raised `AssertionError: eltype(A) == eltype(B)` in `mat_tensor_mul`. `_mul` checks no element type, because GEMM promotes. This closes known issue K7, which leaves `KNOWN_ISSUES.md`. The test is `test/integration/mixed_precision_training.jl`. -- **A `DataLoader`'s two output fields are concrete, its minibatch index set is concrete and draws every sample, and a loader of one sample batches.** `output_dim` and `output_time_steps` were `Union{Int, Nothing}` whatever the loader held, which made a kernel launch read a value of no single type (issue #211); the copy loops of a solution and an ensemble build a 1-based tensor, and the batch indices are one host-built `2 × n` matrix copied to the backend once, with no `@views` (issue #212, which asked for their removal). A `DataLoader` of one sample drew no minibatch (known issue K8) and a loader of matrices never drew its last sample (known issue K12); both now draw every pair, and the count of a trajectory is `cld(n · n_params, batch_size)` in the two files, so a seeded run draws the minibatches it drew before, except the ones the last sample joins. The `JET.report_opt` line of `test/quality/jet.jl` for a launcher that reads a `Union{Int, Nothing}` field is a `@test` rather than a `@test_broken`: the batch half of issue #325 is fixed here. This closes known issues B10, C17, K8 and K12, which leave `KNOWN_ISSUES.md`. B10's `DataLoader(::EnsembleSolution{T, T1, Vector{ST}})` for a `dataser` of `(t, q, q̇)` is reachable — a rigid-body `odeensemble` produces one, and the tutorial and three scripts call it — and `test/data_loader/data_loader.jl` tests it, so the method stays. +- **A `DataLoader`'s two output fields are concrete, its minibatch index set is concrete and draws every sample, and a loader of one sample batches.** `output_dim` and `output_time_steps` were `Union{Int, Nothing}` whatever the loader held, which made a kernel launch read a value of no single type (issue #211); the copy loops of a solution and an ensemble build a 1-based tensor, and the batch indices are one host-built `2 × n` matrix copied to the backend once, with no `@views` (issue #212, which asked for `@views`). A `DataLoader` of one sample drew no minibatch (known issue K8) and a loader of matrices never drew its last sample (known issue K12); both now draw every pair, and the count of a trajectory is `cld(n · n_params, batch_size)` in the two files, so a seeded run draws the minibatches it drew before, except the ones the last sample joins. The `JET.report_opt` line of `test/quality/jet.jl` for a launcher that reads a `Union{Int, Nothing}` field is a `@test` rather than a `@test_broken`: the batch half of issue #325 is fixed here. This closes known issues B10, C17, K8 and K12, which leave `KNOWN_ISSUES.md`. B10's `DataLoader(::EnsembleSolution{T, T1, Vector{ST}})` for a `dataser` of `(t, q, q̇)` is reachable — a rigid-body `odeensemble` produces one, and the tutorial and three scripts call it — and `test/data_loader/data_loader.jl` tests it, so the method stays. - **`assign_output_estimate` is a crop of the last `prediction_window` columns, and its gradient is Zygote's `getindex` adjoint.** Its kernel, `augment_zeros` and the `rrule` are gone: the adjoint is zero on the cropped columns, runs on a device with scalar indexing off, and gives a `Float32` cotangent for a `Float64` seed on `Float32` data, where `augment_zeros` returned the seed's type. The docstring's formula starts at column `T − pw + 1`, which is what the code does. diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index 6158d6831..4b92927db 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -297,6 +297,22 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. are in `test/devices/metal.jl`, and `origin/main` at 97bf0d8 has no `test/metal/` directory. The fix is a change of that path under `src/`. +### K18 · The tensor-sampling formula in the `DataLoader` manual is one short + +- **location:** `docs/src/data_loader/data_loader.md:141` +- **kind:** docs +- **found:** 2026-10-16 +- **evidence:** + Item 1 of *Sampling from a tensor* gives the start positions as + `shuffle(1:(input_time_steps − seq_length − prediction_window))`, and item 4 takes `n_batches` + from them. The code takes + `n = min(input_time_steps − seq_length + 1, target_steps − prediction_window + 1)` start + positions of one parameter index, and counts `cld(n · n_params, batch_size)` minibatches + (`src/data_loader/batch.jl`: `number_of_start_positions` and `number_of_batches`). The formula + is one short of the code's, and the line is byte for byte the same at `d751ecd9` and on this + branch, so no part of the 0.9 work introduced the difference. The correction is a change of + that text. + ## The test suite ### K3 · `test/quality/jet.jl` does not see an unstable kernel-launch argument, nor two barriers in the Poisson-tensor kernel diff --git a/docs/src/data_loader/data_loader.md b/docs/src/data_loader/data_loader.md index 58b8f72a1..86251f7c8 100644 --- a/docs/src/data_loader/data_loader.md +++ b/docs/src/data_loader/data_loader.md @@ -106,7 +106,7 @@ batch = Batch(3) batch(dl) ``` -In the first of those two examples the `autoencoder` keyword was set to `true`, which is why the first index was always `1`. The second sets it to `false`, the default: +In those two examples the `autoencoder` keyword was set to `true`, which is why the first index is always `1`. This changes with `autoencoder = false`, the default: ```@example using GeometricMachineLearning # hide diff --git a/src/data_loader/batch.jl b/src/data_loader/batch.jl index 4a27344be..b1105040b 100644 --- a/src/data_loader/batch.jl +++ b/src/data_loader/batch.jl @@ -99,7 +99,7 @@ function Batch(batch_size::Int, seq_length::Int, prediction_window::Int = seq_le end # The classifier's batch: one draw per sample, with the whole input sequence and the whole output. -# Both lengths are the loader's. Its old prediction window of `0` would gather an output of width 0. +# Its fifth type parameter is `Int`, because both lengths come from a loader that holds an output. function Batch(batch_size::Int, dl::DataLoader{T, AT, OT, K, Int}) where {T, AT, OT <: AbstractArray, K} Batch(batch_size, dl.input_time_steps, dl.output_time_steps) @@ -111,9 +111,8 @@ end Compute the number of batches. One epoch draws every start position of every parameter index once, in minibatches of the batch -size, so the count is `cld(n · n_params, batch_size)` with `n` the number of those positions and -`samples without replacement`: two indices are never drawn twice. The last minibatch holds what is -left. +size, so the count is `cld(n · n_params, batch_size)` with `n` the number of start positions of one +parameter index. No two indices are drawn twice, and the last minibatch holds what is left. # Examples @@ -168,14 +167,16 @@ function number_of_start_positions(dl::DataLoader{T, AT, OT, K}, check_window(min(input_window, target_window), dl, batch) end -# A loader and a batch that leave no start position are a caller's mistake: an epoch of no minibatch -# trains nothing and says nothing, so the three numbers that decide it are named here. +# A loader or a batch that leaves no minibatch is a caller's mistake: an epoch of none trains +# nothing and says nothing, so the numbers that decide it are named here. The message is built only +# where it is thrown, so the check itself allocates nothing. function check_window(n::Int, dl::DataLoader, batch::Batch) - n ≥ 1 || - throw(ArgumentError("the `DataLoader` has $(dl.input_time_steps) time steps, and the " * - "`Batch` asks for a sequence length of $(batch.seq_length) and a " * - "prediction window of $(batch.prediction_window), which leaves no " * - "start position.")) + n ≥ 1 && dl.n_params ≥ 1 || throw(ArgumentError( + "the `DataLoader` has $(dl.input_time_steps) time steps" * + (dl.output === nothing ? "" : " and $(dl.output_time_steps) in its output") * + ", with `n_params = $(dl.n_params)`, and the `Batch` asks for a sequence length of " * + "$(batch.seq_length) and a prediction window of $(batch.prediction_window), which " * + "leaves no minibatch.")) n end diff --git a/src/data_loader/data_loader.jl b/src/data_loader/data_loader.jl index 33b26f9cd..522d62ba2 100644 --- a/src/data_loader/data_loader.jl +++ b/src/data_loader/data_loader.jl @@ -49,12 +49,13 @@ end # The fifth parameter is the type of the two output fields, which is `Nothing` for a loader without # an output and `Int` with one. A caller that names the first four — `DataLoader{T, AT, OT, -# :TimeSeries}(…)`, which is how `GMLDatasets` builds a loader — gets it filled in from the fields. -function DataLoader{T, AT, OT, K}(input::AT, output::OT, input_dim::Int, - input_time_steps::Int, n_params::Int, output_dim, output_time_steps) where { - T, AT, OT, K} - DataLoader{T, AT, OT, K, typeof(output_dim)}(input, output, input_dim, - input_time_steps, n_params, output_dim, output_time_steps) +# :TimeSeries}(…)`, which is how `GMLDatasets` builds a loader — gets it filled in from the fields, +# and the five counts are converted to `Int` whatever integer type the caller passes. +function DataLoader{T, AT, OT, K}(input::AT, output::OT, input_dim, input_time_steps, + n_params, output_dim, output_time_steps) where {T, AT, OT, K} + S = output_dim === nothing ? Nothing : Int + DataLoader{T, AT, OT, K, S}(input, output, input_dim, input_time_steps, n_params, + output_dim, output_time_steps) end # The two decisions `autoencoder` makes at every constructor that takes a data set: the kind it @@ -243,24 +244,29 @@ function DataLoader(data::NamedTuple{(:q, :p), Tuple{VT, VT}}; autoencoder = autoencoder, suppress_info = suppress_info) end -# One 1-based data tensor out of the coordinate of one solution. The `q` of a solution carries the +# The 1-based tensor of one coordinate of one solution, or of each member of an ensemble along the +# parameter axis, with the axes `(system dimension, time steps, parameters)`. +function states_tensor(states::AbstractVector, n_params::Int = 1) + Array{eltype(first(states))}(undef, length(first(states)), length(states), n_params) +end + +# The states of one coordinate into the column `i` of `data`. The `q` of a solution carries the # `DataSeries` time axis `0:n` and a `DataLoader` tensor is 1-based on all three axes, so the states -# are copied in one at a time. -function data_tensor(q::AbstractVector) - data = Array{eltype(first(q))}(undef, length(first(q)), length(q), 1) - for (t, state) in enumerate(q) - data[:, t, 1] = state +# are copied in one at a time and no whole coordinate is built on the way. +function data_tensor!(data, states::AbstractVector, i::Int) + for (t, state) in enumerate(states) + data[:, t, i] = state end data end +data_tensor(states::AbstractVector) = data_tensor!(states_tensor(states), states, 1) + # The tensor of one coordinate over the members of an ensemble, along the parameter axis. function ensemble_data_tensor(solutions, coordinate::Symbol) - first_data = data_tensor(getproperty(first(solutions), coordinate)) - data = similar(first_data, size(first_data, 1), size(first_data, 2), length(solutions)) - for (i, solution) in pairs(solutions) - data[:, :, i] = data_tensor(getproperty(solution, coordinate)) - end + data = states_tensor(getproperty(first(solutions), coordinate), length(solutions)) + foreach(i -> data_tensor!(data, getproperty(solutions[i], coordinate), i), + eachindex(solutions)) data end @@ -341,10 +347,12 @@ function DataLoader(ensemble_solution::EnsembleSolution{T, T1, Vector{ST}}; autoencoder = autoencoder, suppress_info = suppress_info) end +# `Base.copyto!` carries a host array to any backend; `KernelAbstractions.copyto!` has a method for +# the CPU alone, and a device backend such as `JLArray`'s has none. function map_to_new_backend(input::AbstractArray{T}, backend::KernelAbstractions.Backend) where {T} converted = KernelAbstractions.allocate(backend, T, size(input)...) - KernelAbstractions.copyto!(backend, converted, input) + Base.copyto!(converted, input) converted end diff --git a/test/data_loader/accuracy.jl b/test/data_loader/accuracy.jl index 9ae757417..aa4ef45e7 100644 --- a/test/data_loader/accuracy.jl +++ b/test/data_loader/accuracy.jl @@ -3,9 +3,8 @@ # caller and no test is a function nobody notices breaking, so this file exercises it rather than # the docs entry being dropped. # -# The loader comes from `DataLoader(input, output)`, which is C17's constructor: its two element -# types are free, so a real input and an integer one-hot target are one call. Until M2 the only way -# to build this argument was to name the parametric type, and the file said so. +# The loader comes from `DataLoader(input, output)`, whose two element types are free, so a real +# input and an integer one-hot target are one call. using GeometricMachineLearning using Random @@ -55,9 +54,9 @@ end end end -# Edge 10: the crop of `assign_output_estimate` is the last `prediction_window` columns, and its -# gradient is Zygote's `getindex` adjoint — zero on the cropped columns, and a cotangent of the -# primal's element type whatever the seed's. +# The crop of `assign_output_estimate` is the last `prediction_window` columns, and its gradient is +# Zygote's `getindex` adjoint — zero on the cropped columns, and a cotangent of the primal's element +# type whatever the seed's. @testset "assign_output_estimate crops the last columns ($T)" for T in (Float32, Float64) x = T.(reshape(1:(3 * 5 * 4), 3, 5, 4)) @test GeometricMachineLearning.assign_output_estimate(x, 2) == x[:, 4:5, :] diff --git a/test/data_loader/batch.jl b/test/data_loader/batch.jl index 9979f15a7..540a6f1ee 100644 --- a/test/data_loader/batch.jl +++ b/test/data_loader/batch.jl @@ -1,4 +1,4 @@ -# The batching rules of `src/data_loader/batch.jl`. The data are the part's own: ten samples +# The batching rules of `src/data_loader/batch.jl`. The data are ten samples # `x = T.(reshape(1:20, 2, 10))` with `y = 3 .* x`, so that `output == 3 .* input` identifies the # pair of one sample and not the pair of another, and a drawn index pair names the sample it draws. @@ -42,6 +42,7 @@ end @test all(size(input)[1:2] == (2, 1) for input in inputs) for (input, output) in data @test eltype(input) === T + @test eltype(output) === T @test size(output) == size(input) @test output == 3 .* input end @@ -61,20 +62,31 @@ end @test (dl.input_dim, dl.input_time_steps, dl.n_params) == (4, 1, 10) batch = Batch(3) - @test sort(vcat(batch(dl)...)) == [(1, p) for p in 1:10] - minibatches = gathered(dl, batch) + drawn = batch(dl) + @test sort(vcat(drawn...)) == [(1, p) for p in 1:10] + minibatches = map( + indices -> convert_input_and_batch_indices_to_array(dl, batch, indices), + drawn) @test length(minibatches) == 4 @test sort([size(mb.q, 3) for mb in minibatches]) == [1, 3, 3, 3] - for mb in minibatches + for (indices, mb) in zip(drawn, minibatches) @test mb isa NamedTuple{(:q, :p)} @test eltype(mb.q) === T + @test eltype(mb.p) === T @test mb.p == 3 .* mb.q + # each slot holds the sample its own drawn index names, so a `q` and a `p` read from the + # same other column shows here + @test all(mb.q[:, 1, k] == x[:, indices[k][2]] for k in axes(mb.q, 3)) end + # the gathered columns are the columns of `x`, in some order + @test sort(eachcol(hcat((mb.q[:, 1, k] + for mb in minibatches for k in axes(mb.q, 3))...))) == + sort(eachcol(Matrix(x))) end -# K8: one sample is a loader the training loop can draw from. -@testset "one sample is one minibatch" begin - x = Float32.(reshape(1:2, 2, 1)) +# One sample is a loader the training loop can draw from. +@testset "one sample is one minibatch ($T)" for T in (Float32, Float64) + x = T.(reshape(1:2, 2, 1)) dl = DataLoader(x, x; suppress_info = true) batch = Batch(1) minibatches = batch(dl) @@ -82,11 +94,13 @@ end @test only(minibatches) == [(1, 1)] input, output = only(gathered(dl, batch)) @test size(input) == (2, 1, 1) + @test eltype(input) === T + @test eltype(output) === T @test output == input end -# Edge 3: one count, `cld(n · n_params, batch_size)` with `n` the number of start positions, and an -# `ArgumentError` where no window fits — instead of the assertion and the empty epoch before it. +# One count, `cld(n · n_params, batch_size)` with `n` the number of start positions, and an +# `ArgumentError` where no window fits or the parameter axis is empty. @testset "one count, and a window that fits no start position" begin trajectory = DataLoader(rand(Float64, 2, 10); suppress_info = true) @test kind(trajectory) === :TimeSeries @@ -112,10 +126,26 @@ end @test occursin("$(trajectory.input_time_steps)", raised.msg) @test occursin("11", raised.msg) @test occursin("5", raised.msg) + + # a loader whose parameter axis is empty draws no minibatch either + @test_throws ArgumentError Batch(1)(DataLoader(zeros(Float64, 2, 0); + autoencoder = true, suppress_info = true)) + @test_throws ArgumentError Batch(1)(DataLoader(zeros(Float64, 2, 3, 0); suppress_info = true)) + + # where the output's time steps bind the window, the message names them + with_output = DataLoader(rand(Float64, 2, 10, 3), rand(Float64, 2, 1, 3); + suppress_info = true) + raised = try + Batch(2, 1, 2)(with_output) + catch e + e + end + @test raised isa ArgumentError + @test occursin("10 time steps and 1 in its output", raised.msg) end -# Edges 4 and 5: the classifier's loader, and `Batch(batch_size, dl)` — one draw per sample with the -# whole sequence and the whole output, which is what GMLDatasets calls. +# The classifier's loader, and `Batch(batch_size, dl)` — one draw per sample with the whole +# sequence and the whole output, which is what GMLDatasets calls. @testset "Batch(batch_size, dl) draws every sample with the whole sequence ($T)" for T in (Float32, Float64) n_classes, n_patches, n_images = 10, 4, 10 input = rand(T, 9, n_patches, n_images) @@ -143,17 +173,17 @@ end end end -# Edge 2: the target is the output when there is one, whatever kind the loader is typed as. A +# The target is the output when there is one, whatever kind the loader is typed as. A # `DataLoader{T, A3, Array{Int, 3}, :TimeSeries}` is what GMLDatasets and two tests build by naming # the type, and it must batch like the `:RegularData` loader of the same data. -@testset "a loader with an output is samples whatever its kind" begin - input = rand(Float32, 9, 4, 10) +@testset "a loader with an output is samples whatever its kind ($T)" for T in (Float32, Float64) + input = rand(T, 9, 4, 10) output = zeros(Int, 10, 1, 10) for (k, c) in pairs(rand(1:10, 10)) output[c, 1, k] = 1 end dl = DataLoader(input, output; suppress_info = true) - dl_time_series = DataLoader{Float32, typeof(input), typeof(output), :TimeSeries}( + dl_time_series = DataLoader{T, typeof(input), typeof(output), :TimeSeries}( input, output, 9, 4, 10, 10, 1) for batch in (Batch(2, dl), Batch(3, 1, 1)) @@ -163,8 +193,10 @@ end reference = batch(dl) @test asked == reference for (i, j) in zip(asked, reference) - @test convert_input_and_batch_indices_to_array(dl_time_series, batch, i) == - convert_input_and_batch_indices_to_array(dl, batch, j) + here = convert_input_and_batch_indices_to_array(dl_time_series, batch, i) + @test here == convert_input_and_batch_indices_to_array(dl, batch, j) + @test eltype(here[1]) === T + @test eltype(here[2]) === Int end end end diff --git a/test/data_loader/data_loader.jl b/test/data_loader/data_loader.jl index f86c79b10..2c5df8627 100644 --- a/test/data_loader/data_loader.jl +++ b/test/data_loader/data_loader.jl @@ -1,8 +1,7 @@ # The constructors of `src/data_loader/data_loader.jl`: the `autoencoder` default and the keyword -# set of every method (the part's check 4), the fifth type parameter of the struct (#211), the copy -# loops out of a solution and an ensemble (B10, edge 11), the element-type conversion of -# `DataLoader(dl, T)`, and the `(input, output)` method that lets a classifier's data be built -# without naming the parametric type (C17). +# set of every method, the fifth type parameter of the struct (#211), the copy loops out of a +# solution and an ensemble, the element-type conversion of `DataLoader(dl, T)`, and the +# `(input, output)` method that lets a classifier's data be built without naming the parametric type. using GeometricMachineLearning using Test @@ -58,9 +57,9 @@ end @test_throws MethodError DataLoader(ensemble, true) end -# B10 and edge 11: the copy loops of `data_tensors_from_geometric_solution` build one 1-based tensor -# per coordinate. The rigid-body `odeensemble` is the case whose `dataser` carries `(t, q, q̇)` and -# no `p`, which is the method B10 called unreachable. +# The copy loops of `data_tensors_from_geometric_solution` build one 1-based tensor per coordinate. +# The rigid-body `odeensemble` is the case whose `dataser` carries `(t, q, q̇)` and no `p`, which is +# the case its own method is for. @testset "a solution and an ensemble become 1-based tensors" begin ics = [(q = [sin(v), 0.0, cos(v)],) for v in 0.1:0.5:2.0] ensemble = integrate( @@ -96,6 +95,17 @@ end @test dl_qp.input.q[:, t + 1, i] == sol.q[t] @test dl_qp.input.p[:, t + 1, i] == sol.p[t] end + + # and one `(q, p)` solution: the two coordinates are not swapped + solution_qp = integrate( + hodeproblem(; timespan = (0.0, 1.0), timestep = 0.25), ImplicitMidpoint()) + dl_qp_one = DataLoader(solution_qp; suppress_info = true) + @test dl_qp_one.input isa NamedTuple{(:q, :p)} + @test size(dl_qp_one.input.q) == (1, length(solution_qp.t), 1) + for t in 0:(length(solution_qp.t) - 1) + @test dl_qp_one.input.q[:, t + 1, 1] == solution_qp.q[t] + @test dl_qp_one.input.p[:, t + 1, 1] == solution_qp.p[t] + end end # The reproduction script `scripts/reproduction/symplectic_autoencoders/training.jl` has a `solve!` @@ -122,6 +132,25 @@ end end end +# The four-parameter constructor fills the fifth type parameter in, and it takes the counts as any +# integer type: `GMLDatasets` builds a loader by naming the first four and handing it the `size`s. +@testset "the four-parameter constructor converts its counts" begin + input = rand(Float32, 2, 3, 4) + output = zeros(Int, 3, 1, 4) + dl = DataLoader{Float32, typeof(input), typeof(output), :TimeSeries}(input, output, + Int32(2), Int32(3), Int32(4), Int32(3), Int32(1)) + @test dl isa DataLoader{Float32, Array{Float32, 3}, Array{Int, 3}, :TimeSeries, Int} + @test (dl.input_dim, dl.input_time_steps, dl.n_params) == (2, 3, 4) + @test (dl.output_dim, dl.output_time_steps) == (3, 1) + @test dl.output === output + + no_output = DataLoader{Float32, typeof(input), Nothing, :RegularData}(input, nothing, + Int8(2), Int8(3), Int8(4), nothing, nothing) + @test no_output isa + DataLoader{Float32, Array{Float32, 3}, Nothing, :RegularData, Nothing} + @test (no_output.output_dim, no_output.output_time_steps) == (nothing, nothing) +end + # The conversion to another element type keeps the kind and the counts, and the other backend # branch is the identity. @testset "DataLoader(dl, T) converts and keeps the kind ($T)" for T in (Float32, Float64) @@ -139,8 +168,8 @@ end @test DataLoader(dl, other).input === dl.input end -# C17: a classifier's real input and integer one-hot output, without naming the parametric type. -@testset "C17: a loader of an input and an integer output" begin +# A classifier's real input and integer one-hot output, without naming the parametric type. +@testset "a loader of an input and an integer output" begin input = rand(Float32, 4, 3, 6) output = zeros(Int64, 3, 1, 6) dl = DataLoader(input, output; suppress_info = true) From 07ab13e900b32e8746bfbc0a4e38595a6384e7f1 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 16 Oct 2026 12:46:35 +0200 Subject: [PATCH 5/8] DataLoader: the accuracy header's precision claim --- test/data_loader/accuracy.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/test/data_loader/accuracy.jl b/test/data_loader/accuracy.jl index aa4ef45e7..8da9c4f7b 100644 --- a/test/data_loader/accuracy.jl +++ b/test/data_loader/accuracy.jl @@ -1,10 +1,10 @@ # `accuracy` is documented API -- `docs/src/data_loader/data_loader.md` renders both its methods -- -# and nothing in the package, the tests or the scripts called it. A documented function with no -# caller and no test is a function nobody notices breaking, so this file exercises it rather than -# the docs entry being dropped. +# and no other function of the package and no script calls it. A documented function with no +# caller and no test is a function nobody notices breaking, so this file exercises it. # -# The loader comes from `DataLoader(input, output)`, whose two element types are free, so a real -# input and an integer one-hot target are one call. +# The loader comes from `DataLoader(input, output)`. The output's element type is the input's +# or an integer, so a real input and an integer one-hot target are one call; any other pair of +# element types is a `MethodError`. using GeometricMachineLearning using Random From 62826284c2c047e71a1d683c2d3958e796b0c626 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 16 Oct 2026 13:40:02 +0200 Subject: [PATCH 6/8] DataLoader: the pre-PR verifier's fixes and the constructor's inference The two tensor constructors keep data_set_kind(autoencoder) and gain Base.@constprop :aggressive, which restores the concrete kind the base inferred, without the duplicated branches that took the three-file code-line count to 284. The count is 274 of the part's 275. Six comments now state the current contracts, KNOWN_ISSUES.md gains K18 for the manual's sampling formula, and the CHANGELOG entry no longer claims the autoencoder default and suppress_info for the copying constructor. --- CHANGELOG.md | 2 +- KNOWN_ISSUES.md | 16 +++++++--------- src/data_loader/batch.jl | 2 +- src/data_loader/data_loader.jl | 4 ++-- src/data_loader/tensor_assign.jl | 4 ++-- test/data_loader/autoencoder_keyword.jl | 16 ++++------------ test/data_loader/data_loader.jl | 8 ++++++++ test/data_loader/optimizer_functor_with_adam.jl | 2 +- 8 files changed, 26 insertions(+), 28 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 029b48376..1c72e953c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -41,7 +41,7 @@ breaking release). ### Changed (breaking) -- **`DataLoader`'s `autoencoder` default is `false` in every constructor, and a loader given an input and an output is a set of samples.** A matrix or a vector is a time series unless `autoencoder = true` is passed: that keyword is what reads a matrix's second axis as the parameter index, and the matrix and the vector method defaulted to `true`, so migrate `DataLoader(matrix)` to `DataLoader(matrix; autoencoder = true)` where the columns are samples — `scripts/reproduction/symplectic_autoencoders/training.jl` and the tutorials do. `DataLoader(input, output)` builds a `:RegularData` loader whatever its kind, so its target is its output and never the input shifted by the sequence length: at the base a `DataLoader(x, y)` of ten pairs drew the samples 1 to 9 and, where the input and the output differed in array type, returned one minibatch holding the whole data set. The count is `cld(n · n_params, batch_size)` with `n` the start positions of one parameter index, and a loader and a `Batch` that leave no start position raise an `ArgumentError` naming the time steps, the sequence length and the prediction window, where the base asserted and an epoch could be empty. `Batch(batch_size, dl)` takes its two lengths from the loader, so a classifier's minibatch holds the whole patch sequence and the whole one-hot output of each of its samples. `suppress_info` is a keyword in every method, which makes `DataLoader(solution, true)` a `MethodError`. +- **`DataLoader`'s `autoencoder` defaults to `false` in every constructor that takes it, except the copying constructor, where `autoencoder = nothing` inherits the loader's kind, and a loader given an input and an output is a set of samples.** A matrix or a vector is a time series unless `autoencoder = true` is passed: that keyword is what reads a matrix's second axis as the parameter index, and the matrix and the vector method defaulted to `true`, so migrate `DataLoader(matrix)` to `DataLoader(matrix; autoencoder = true)` where the columns are samples — `scripts/reproduction/symplectic_autoencoders/training.jl` and the tutorials do. `DataLoader(input, output)` builds a `:RegularData` loader whatever its kind, so its target is its output and never the input shifted by the sequence length: at the base a `DataLoader(x, y)` of ten pairs drew the samples 1 to 9 and, where the input and the output differed in array type, returned one minibatch holding the whole data set. The count is `cld(n · n_params, batch_size)` with `n` the start positions of one parameter index, and a loader and a `Batch` that leave no start position raise an `ArgumentError` naming the time steps, the sequence length and the prediction window, where the base asserted and an epoch could be empty. `Batch(batch_size, dl)` takes its two lengths from the loader, so a classifier's minibatch holds the whole patch sequence and the whole one-hot output of each of its samples. `suppress_info` is a keyword of every constructor that takes a data set, which makes `DataLoader(solution, true)` a `MethodError`. - **The two output fields of `DataLoader` carry a fifth type parameter.** `output_dim` and `output_time_steps` are the type `S <: Union{Int, Nothing}` of the loader, so `DataLoader{T, AT, OT, :TimeSeries}` still names a loader and a *full* type name gains `, Nothing` for a loader without an output; a four-parameter constructor `DataLoader{T, AT, OT, K}(input, output, input_dim, input_time_steps, n_params, output_dim, output_time_steps)` fills `S` in from the fields. diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index 4b92927db..bc059850d 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -299,19 +299,17 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. ### K18 · The tensor-sampling formula in the `DataLoader` manual is one short -- **location:** `docs/src/data_loader/data_loader.md:141` +- **location:** `docs/src/data_loader/data_loader.md:143` - **kind:** docs - **found:** 2026-10-16 - **evidence:** Item 1 of *Sampling from a tensor* gives the start positions as - `shuffle(1:(input_time_steps − seq_length − prediction_window))`, and item 4 takes `n_batches` - from them. The code takes - `n = min(input_time_steps − seq_length + 1, target_steps − prediction_window + 1)` start - positions of one parameter index, and counts `cld(n · n_params, batch_size)` minibatches - (`src/data_loader/batch.jl`: `number_of_start_positions` and `number_of_batches`). The formula - is one short of the code's, and the line is byte for byte the same at `d751ecd9` and on this - branch, so no part of the 0.9 work introduced the difference. The correction is a change of - that text. + `shuffle(1:(input_time_steps − seq_length − prediction_window))`. For the time-series loader + without an output shown on that page, `number_of_start_positions` takes + `n = input_time_steps − seq_length − prediction_window + 1`, and `number_of_batches` counts + `cld(n · n_params, batch_size)` minibatches (`src/data_loader/batch.jl`). For the example's + five time steps and `Batch(4, 2, 0)`, the formula gives three start positions per parameter + and the code gives four. The correction is to add `+ 1` to the formula's upper bound. ## The test suite diff --git a/src/data_loader/batch.jl b/src/data_loader/batch.jl index b1105040b..836a9941e 100644 --- a/src/data_loader/batch.jl +++ b/src/data_loader/batch.jl @@ -99,7 +99,7 @@ function Batch(batch_size::Int, seq_length::Int, prediction_window::Int = seq_le end # The classifier's batch: one draw per sample, with the whole input sequence and the whole output. -# Its fifth type parameter is `Int`, because both lengths come from a loader that holds an output. +# The loader's fifth type parameter is `Int`, because both lengths come from a loader with an output. function Batch(batch_size::Int, dl::DataLoader{T, AT, OT, K, Int}) where {T, AT, OT <: AbstractArray, K} Batch(batch_size, dl.input_time_steps, dl.output_time_steps) diff --git a/src/data_loader/data_loader.jl b/src/data_loader/data_loader.jl index 522d62ba2..226daef17 100644 --- a/src/data_loader/data_loader.jl +++ b/src/data_loader/data_loader.jl @@ -120,7 +120,7 @@ DataLoader{Int64, Array{Int64, 3}, Nothing, :TimeSeries, Nothing}([1 4; 2 5; 3 6 ``` the `@info` statement is not printed. """ -function DataLoader( +Base.@constprop :aggressive function DataLoader( data::AbstractArray{T, 3}; autoencoder::Bool = false, suppress_info = false) where {T} if !suppress_info @info "You have provided a tensor with three axes as input. They will be interpreted as \n (i) system dimension, (ii) number of time steps and (iii) number of params." @@ -225,7 +225,7 @@ function DataLoader(data::NamedTuple{(:q, :p), Tuple{AT, AT}}; autoencoder = autoencoder, suppress_info = true) end -function DataLoader(data::NamedTuple{(:q, :p), Tuple{AT, AT}}; +Base.@constprop :aggressive function DataLoader(data::NamedTuple{(:q, :p), Tuple{AT, AT}}; autoencoder::Bool = false, suppress_info = false) where {T, AT <: AbstractArray{T, 3}} if !suppress_info diff --git a/src/data_loader/tensor_assign.jl b/src/data_loader/tensor_assign.jl index fe76ff5dc..5f6d9570a 100644 --- a/src/data_loader/tensor_assign.jl +++ b/src/data_loader/tensor_assign.jl @@ -1,6 +1,6 @@ # Assigning a batch and its target out of the full data tensor is `src/data_loader/batch.jl`'s job: -# its gather kernel indexes by a vector of `(time_step, parameter)` tuples. What is here estimates -# the output instead. +# its gather kernel reads a matrix of `(time_step, parameter)` indices. This file crops the +# estimated output instead. @doc raw""" assign_output_estimate(full_output, prediction_window) diff --git a/test/data_loader/autoencoder_keyword.jl b/test/data_loader/autoencoder_keyword.jl index 408dd2a70..4e9b6dc51 100644 --- a/test/data_loader/autoencoder_keyword.jl +++ b/test/data_loader/autoencoder_keyword.jl @@ -1,10 +1,5 @@ -# Each of these constructors branches on `autoencoder`, so a value that is neither `true` nor -# `false` -- `nothing` being the case that actually arises, since it is the sentinel a caller -# reaches for when in doubt -- must be rejected at the keyword rather than inside the branch. -# Declaring the keyword `autoencoder::Bool` makes it a `TypeError` at the call the keyword is -# passed to, which is where the mistake is. Without the declaration such a value falls through an -# `if`/`elseif` with no `else`, and the constructor returns `nothing` instead of a `DataLoader` -# with no indication of why. +# Data-set constructors accept `autoencoder::Bool`. A value outside `Bool`, such as `nothing`, +# is rejected at the keyword with a `TypeError` rather than interpreted as a loader kind. using GeometricMachineLearning using Test @@ -35,11 +30,8 @@ end @test_throws TypeError DataLoader(data; autoencoder = nothing, suppress_info = true) end -# The fifth site, `DataLoader(dl, backend)`, is not one of the four above: its `autoencoder = -# nothing` default is not a mistake to reject, it is the sentinel meaning "inherit from `dl`". It -# takes `Union{Nothing, Bool}` instead, and its `DT = if …` block carries the `else` that the type -# gives the other four, so a value outside that union is a `TypeError` there too and nothing falls -# through. +# `DataLoader(dl, backend)` accepts `autoencoder::Union{Nothing, Bool}`: `nothing` inherits the +# loader's kind, a Boolean selects it, and a value outside the union raises a `TypeError`. @testset "DataLoader(dl, backend): nothing means inherit, and nothing falls through" begin dl = DataLoader(rand(Float32, 2, 3, 4); autoencoder = true, suppress_info = true) diff --git a/test/data_loader/data_loader.jl b/test/data_loader/data_loader.jl index 2c5df8627..88c311512 100644 --- a/test/data_loader/data_loader.jl +++ b/test/data_loader/data_loader.jl @@ -23,6 +23,14 @@ kind(dl::DataLoader) = typeof(dl).parameters[4] @test kind(DataLoader(rand(T, 10); suppress_info = true)) === :TimeSeries @test kind(DataLoader(tensor; suppress_info = true)) === :TimeSeries @test kind(DataLoader((q = matrix, p = matrix); suppress_info = true)) === :TimeSeries + # A literal keyword fixes the loader's kind, so the tensor constructors infer a concrete type. + @test (@inferred (() -> DataLoader(tensor; suppress_info = true))()) isa DataLoader + @test (@inferred (() -> DataLoader((q = tensor, p = tensor); + suppress_info = true))()) isa DataLoader + @test (@inferred (() -> DataLoader(tensor; autoencoder = true, + suppress_info = true))()) isa DataLoader + @test (@inferred (() -> DataLoader((q = tensor, p = tensor); autoencoder = true, + suppress_info = true))()) isa DataLoader @test kind(DataLoader((q = tensor, p = tensor); suppress_info = true)) === :TimeSeries @test kind(DataLoader((q = rand(T, 10), p = rand(T, 10)); suppress_info = true)) === :TimeSeries diff --git a/test/data_loader/optimizer_functor_with_adam.jl b/test/data_loader/optimizer_functor_with_adam.jl index 216a286a9..14a571375 100644 --- a/test/data_loader/optimizer_functor_with_adam.jl +++ b/test/data_loader/optimizer_functor_with_adam.jl @@ -35,7 +35,7 @@ function test_optimization_with_adam(; T = Float32, input_dim = 9, n_patches = 4 T = T, input_dim = input_dim, n_patches = n_patches, n_images = n_images, n_classes = n_classes) - # The batch is drawn over the patches and the images: `Batch(2, dl)` takes the loader's own + # The batch is drawn over the images: `Batch(2, dl)` takes the loader's own # sequence length and output width, so one minibatch holds the whole patch sequence and the # whole one-hot output of its two images. batch = Batch(2, dl) From 6342d909e3bc7f9b2454f817ceda5ce94c024574 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 16 Oct 2026 15:29:35 +0200 Subject: [PATCH 7/8] Review fixes for part M2: one gather signature, one copy loop, concrete matrix loaders - gather takes its sizes from the arrays it reads: the first axis from the data, the batch size from the index matrix. The dim, batch_size and the halving for a (q, p) pair go. - The gather kernel is gather_kernel!, and the batch functor tests the output by one form, OT <: AbstractArray, as number_of_start_positions does. - data_tensor builds the tensor of one coordinate over a vector of trajectories; states_tensor, data_tensor! and ensemble_data_tensor go. - map_to_new_backend returns an array already on the backend, so a CPU minibatch no longer copies its index matrix, and the copying constructor drops its own backend branch. - The matrix and (q, p)-matrix constructors carry Base.@constprop :aggressive, so the matrix and vector constructors infer one concrete loader for a literal autoencoder; the tests assert it. - tensor_assign.jl loses the file-top comment that described batch.jl. - CHANGELOG: "the base" is 0.8.0, Batch(batch_size) on a classifier's loader gets its migration, and the two-element-type claim names the 3-tensor method. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 6 ++-- src/data_loader/batch.jl | 34 ++++++++---------- src/data_loader/data_loader.jl | 59 +++++++++++++------------------- src/data_loader/tensor_assign.jl | 4 --- test/data_loader/data_loader.jl | 9 +++++ test/quality/jet.jl | 4 +-- 6 files changed, 52 insertions(+), 64 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1c72e953c..d2e37fa4a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -41,7 +41,7 @@ breaking release). ### Changed (breaking) -- **`DataLoader`'s `autoencoder` defaults to `false` in every constructor that takes it, except the copying constructor, where `autoencoder = nothing` inherits the loader's kind, and a loader given an input and an output is a set of samples.** A matrix or a vector is a time series unless `autoencoder = true` is passed: that keyword is what reads a matrix's second axis as the parameter index, and the matrix and the vector method defaulted to `true`, so migrate `DataLoader(matrix)` to `DataLoader(matrix; autoencoder = true)` where the columns are samples — `scripts/reproduction/symplectic_autoencoders/training.jl` and the tutorials do. `DataLoader(input, output)` builds a `:RegularData` loader whatever its kind, so its target is its output and never the input shifted by the sequence length: at the base a `DataLoader(x, y)` of ten pairs drew the samples 1 to 9 and, where the input and the output differed in array type, returned one minibatch holding the whole data set. The count is `cld(n · n_params, batch_size)` with `n` the start positions of one parameter index, and a loader and a `Batch` that leave no start position raise an `ArgumentError` naming the time steps, the sequence length and the prediction window, where the base asserted and an epoch could be empty. `Batch(batch_size, dl)` takes its two lengths from the loader, so a classifier's minibatch holds the whole patch sequence and the whole one-hot output of each of its samples. `suppress_info` is a keyword of every constructor that takes a data set, which makes `DataLoader(solution, true)` a `MethodError`. +- **`DataLoader`'s `autoencoder` defaults to `false` in every constructor that takes it, except the copying constructor, where `autoencoder = nothing` inherits the loader's kind, and a loader given an input and an output is a set of samples.** A matrix or a vector is a time series unless `autoencoder = true` is passed: that keyword is what reads a matrix's second axis as the parameter index, and the matrix and the vector method defaulted to `true`, so migrate `DataLoader(matrix)` to `DataLoader(matrix; autoencoder = true)` where the columns are samples — `scripts/reproduction/symplectic_autoencoders/training.jl` and the tutorials do. `DataLoader(input, output)` builds a `:RegularData` loader whatever its kind, so its target is its output and never the input shifted by the sequence length: in 0.8.0 a `DataLoader(x, y)` of ten pairs drew the samples 1 to 9 and, where the input and the output differed in array type, returned one minibatch holding the whole data set. The count is `cld(n · n_params, batch_size)` with `n` the start positions of one parameter index, and a loader and a `Batch` that leave no start position raise an `ArgumentError` naming the time steps, the sequence length and the prediction window, where 0.8.0 asserted and an epoch could be empty. `Batch(batch_size, dl)` takes its two lengths from the loader, so a classifier's minibatch holds the whole patch sequence and the whole one-hot output of each of its samples. On a classifier's loader, `Batch(batch_size)` now draws only the first patch of each sample, where 0.8.0's minibatch held the whole input sequence, and nothing raises; migrate it to `Batch(batch_size, dl)`. `suppress_info` is a keyword of every constructor that takes a data set, which makes `DataLoader(solution, true)` a `MethodError`. - **The two output fields of `DataLoader` carry a fifth type parameter.** `output_dim` and `output_time_steps` are the type `S <: Union{Int, Nothing}` of the loader, so `DataLoader{T, AT, OT, :TimeSeries}` still names a loader and a *full* type name gains `, Nothing` for a loader without an output; a four-parameter constructor `DataLoader{T, AT, OT, K}(input, output, input_dim, input_time_steps, n_params, output_dim, output_time_steps)` fills `S` in from the fields. @@ -77,7 +77,7 @@ breaking release). - **A network with a Stiefel leaf trains on data of the other precision.** Training on data whose element type differs from the network's raised `AssertionError: eltype(A) == eltype(B)` in `mat_tensor_mul`. `_mul` checks no element type, because GEMM promotes. This closes known issue K7, which leaves `KNOWN_ISSUES.md`. The test is `test/integration/mixed_precision_training.jl`. -- **A `DataLoader`'s two output fields are concrete, its minibatch index set is concrete and draws every sample, and a loader of one sample batches.** `output_dim` and `output_time_steps` were `Union{Int, Nothing}` whatever the loader held, which made a kernel launch read a value of no single type (issue #211); the copy loops of a solution and an ensemble build a 1-based tensor, and the batch indices are one host-built `2 × n` matrix copied to the backend once, with no `@views` (issue #212, which asked for `@views`). A `DataLoader` of one sample drew no minibatch (known issue K8) and a loader of matrices never drew its last sample (known issue K12); both now draw every pair, and the count of a trajectory is `cld(n · n_params, batch_size)` in the two files, so a seeded run draws the minibatches it drew before, except the ones the last sample joins. The `JET.report_opt` line of `test/quality/jet.jl` for a launcher that reads a `Union{Int, Nothing}` field is a `@test` rather than a `@test_broken`: the batch half of issue #325 is fixed here. This closes known issues B10, C17, K8 and K12, which leave `KNOWN_ISSUES.md`. B10's `DataLoader(::EnsembleSolution{T, T1, Vector{ST}})` for a `dataser` of `(t, q, q̇)` is reachable — a rigid-body `odeensemble` produces one, and the tutorial and three scripts call it — and `test/data_loader/data_loader.jl` tests it, so the method stays. +- **A `DataLoader`'s two output fields are concrete, its minibatch index set is concrete and draws every sample, and a loader of one sample batches.** `output_dim` and `output_time_steps` were `Union{Int, Nothing}` whatever the loader held, which made a kernel launch read a value of no single type (issue #211); the copies of a solution and of an ensemble both go through one helper, `data_tensor`, which builds a 1-based tensor, and the batch indices are one host-built `2 × n` matrix copied to the backend once, with no `@views` (issue #212, which asked for `@views`). A `DataLoader` of one sample drew no minibatch (known issue K8) and a loader of matrices never drew its last sample (known issue K12); both now draw every pair, and the count of a trajectory is `cld(n · n_params, batch_size)`, so a seeded run draws the minibatches it drew before, except the ones the last sample joins. The `JET.report_opt` line of `test/quality/jet.jl` for a launcher that reads a `Union{Int, Nothing}` field is a `@test` rather than a `@test_broken`: the batch half of issue #325 is fixed here. This closes known issues B10, C17, K8 and K12, which leave `KNOWN_ISSUES.md`. B10's `DataLoader(::EnsembleSolution{T, T1, Vector{ST}})` for a `dataser` of `(t, q, q̇)` is reachable — a rigid-body `odeensemble` produces one, and the tutorial and three scripts call it — and `test/data_loader/data_loader.jl` tests it, so the method stays. - **`assign_output_estimate` is a crop of the last `prediction_window` columns, and its gradient is Zygote's `getindex` adjoint.** Its kernel, `augment_zeros` and the `rrule` are gone: the adjoint is zero on the cropped columns, runs on a device with scalar indexing off, and gives a `Float32` cotangent for a `Float64` seed on `Float32` data, where `augment_zeros` returned the seed's type. The docstring's formula starts at column `T − pw + 1`, which is what the code does. @@ -94,7 +94,7 @@ breaking release). - **`Pkg.test` gains a `JLArray` device test in the `core` group.** `test/devices/jlarray.jl` runs `tensor_tensor_mul` and its pullback, `_mul` with each structured weight in the forward pass and in the storage gradient, and `assign_q_and_p` on a `JLArray` under `allowscalar(false)`. `JLArrays`, `GPUArraysCore` and `ChainRulesCore` join `test/Project.toml`, with `JLArrays = "0.3"`. `NNlib.batched_mul` is not used for `tensor_tensor_mul`, because on a backend with no NNlib extension, JLArray and Metal among them, it falls back to scalar indexing. -- **`DataLoader(input, output)` accepts two element types, one of them an integer.** It is one method, `DataLoader(input::AbstractArray{T, 3}, output::AbstractArray{T1, 3}) where {T, T1 <: Union{T, Integer}}`, so the output's element type is the input's or an integer: a classifier's data — a real input, an integer one-hot target — are one call and `accuracy` is reachable without naming the parametric type (known issue C17, which leaves `KNOWN_ISSUES.md`). Any other pair of element types is a `MethodError`. +- **`DataLoader(input, output)` accepts two element types, one of them an integer, for 3-tensors.** It is one method, `DataLoader(input::AbstractArray{T, 3}, output::AbstractArray{T1, 3}) where {T, T1 <: Union{T, Integer}}`, so the output's element type is the input's or an integer: a classifier's data — a real input, an integer one-hot target — are one call and `accuracy` is reachable without naming the parametric type (known issue C17, which leaves `KNOWN_ISSUES.md`). Any other pair of element types is a `MethodError` for 3-tensors. The matrix and vector `(input, output)` methods take one element type `T` for both arrays, so a `Float32` matrix input with an `Int` matrix output is a `MethodError`. ### Changed diff --git a/src/data_loader/batch.jl b/src/data_loader/batch.jl index 836a9941e..3647f03a5 100644 --- a/src/data_loader/batch.jl +++ b/src/data_loader/batch.jl @@ -192,7 +192,7 @@ function (batch::Batch)(dl::DataLoader) batch_over_two_axes(batch, number_of_start_positions(dl, batch), dl.n_params) end -@kernel function assign_from_vector_of_tuples_kernel!(gathered, data, indices, offset) +@kernel function gather_kernel!(gathered, data, indices, offset) i, j, k = @index(Global, NTuple) gathered[i, j, k] = data[i, indices[1, k] + offset + j - 1, indices[2, k]] @@ -207,21 +207,19 @@ end # The gather kernel, applied to a tensor or to the `q` and `p` of a `(q, p)` pair. The gathered # tensor holds `width` positions from `offset` on, for every drawn start position and parameter -# index, and has the element type of the data it reads — so one integer target stays integer. The -# first axis is the loader's own `output_dim`/`input_dim`, halved for each of a `(q, p)` pair. -function gather(data::AbstractArray{T, 3}, dim::Int, indices::AbstractArray{Int, 2}, - offset::Int, width::Int, batch_size::Int) where {T} +# index, and has the element type of the data it reads — so one integer target stays integer. +function gather(data::AbstractArray{T, 3}, indices::AbstractMatrix{Int}, offset::Int, + width::Int) where {T} backend = networkbackend(data) - gathered = KernelAbstractions.allocate(backend, T, dim, width, batch_size) - gather! = assign_from_vector_of_tuples_kernel!(backend) - gather!(gathered, data, indices, offset, ndrange = size(gathered)) + gathered = KernelAbstractions.allocate( + backend, T, size(data, 1), width, size(indices, 2)) + gather_kernel!(backend)(gathered, data, indices, offset, ndrange = size(gathered)) gathered end -function gather(data::NamedTuple{(:q, :p), Tuple{AT, AT}}, dim::Int, - indices::AbstractArray{Int, 2}, offset::Int, width::Int, - batch_size::Int) where {T, AT <: AbstractArray{T, 3}} - map(tensor -> gather(tensor, dim ÷ 2, indices, offset, width, batch_size), data) +function gather(data::NamedTuple{(:q, :p), Tuple{AT, AT}}, indices::AbstractMatrix{Int}, + offset::Int, width::Int) where {AT <: AbstractArray{<:Any, 3}} + map(tensor -> gather(tensor, indices, offset, width), data) end @doc raw""" @@ -257,18 +255,14 @@ GeometricMachineLearning.convert_input_and_batch_indices_to_array(dl, batch, bat function convert_input_and_batch_indices_to_array( dl::DataLoader{T, AT, OT, K}, batch::Batch, batch_indices_tuple::Vector{Tuple{Int, Int}}) where {T, AT, OT, K} - batch_size = length(batch_indices_tuple) indices = convert_vector_of_tuples_to_matrix(networkbackend(dl), batch_indices_tuple) - input = gather(dl.input, dl.input_dim, indices, 0, batch.seq_length, batch_size) + input = gather(dl.input, indices, 0, batch.seq_length) # a loader that holds an output is a set of samples whatever its kind: the target is its output; # without one, a time series predicts its own continuation and a set of samples has no target - if dl.output !== nothing - input, - gather(dl.output, dl.output_dim, indices, 0, batch.prediction_window, batch_size) + if OT <: AbstractArray + (input, gather(dl.output, indices, 0, batch.prediction_window)) elseif K === :TimeSeries - input, - gather(dl.input, dl.input_dim, indices, batch.seq_length, batch.prediction_window, - batch_size) + (input, gather(dl.input, indices, batch.seq_length, batch.prediction_window)) else input end diff --git a/src/data_loader/data_loader.jl b/src/data_loader/data_loader.jl index 226daef17..f5f402434 100644 --- a/src/data_loader/data_loader.jl +++ b/src/data_loader/data_loader.jl @@ -147,7 +147,7 @@ a tensor of shape `(size(data)..., 1)`: one trajectory of one parameter. With `a the matrix is a set of samples instead, so that axis is the parameter index and the time axis has length one. """ -function DataLoader( +Base.@constprop :aggressive function DataLoader( data::AbstractMatrix{T}; autoencoder::Bool = false, suppress_info = false) where {T} if !suppress_info @info "You have provided a matrix as input. The axes will be interpreted as (i) system dimension and (ii) $(autoencoder ? "number of parameters" : "number of time steps")." @@ -213,7 +213,7 @@ In this case the field `input_dim` of `DataLoader` is interpreted as the sum of Apart from this the input is treated similarly as if it were an `Array`, i.e. everything is converted to tensors internally. See e.g. [`DataLoader{::AbstractArray{<:Number, 3}}`](@ref). """ -function DataLoader(data::NamedTuple{(:q, :p), Tuple{AT, AT}}; +Base.@constprop :aggressive function DataLoader(data::NamedTuple{(:q, :p), Tuple{AT, AT}}; autoencoder::Bool = false, suppress_info = false) where {T, AT <: AbstractMatrix{T}} if !suppress_info @@ -244,42 +244,31 @@ function DataLoader(data::NamedTuple{(:q, :p), Tuple{VT, VT}}; autoencoder = autoencoder, suppress_info = suppress_info) end -# The 1-based tensor of one coordinate of one solution, or of each member of an ensemble along the -# parameter axis, with the axes `(system dimension, time steps, parameters)`. -function states_tensor(states::AbstractVector, n_params::Int = 1) - Array{eltype(first(states))}(undef, length(first(states)), length(states), n_params) -end - -# The states of one coordinate into the column `i` of `data`. The `q` of a solution carries the -# `DataSeries` time axis `0:n` and a `DataLoader` tensor is 1-based on all three axes, so the states -# are copied in one at a time and no whole coordinate is built on the way. -function data_tensor!(data, states::AbstractVector, i::Int) - for (t, state) in enumerate(states) - data[:, t, i] = state +# The 1-based tensor of one coordinate over a vector of trajectories — one solution, or the members +# of an ensemble — with the axes `(system dimension, time steps, trajectories)`. The `q` of a +# solution carries the `DataSeries` time axis `0:n`, so the states are copied in one at a time. +function data_tensor(trajectories::AbstractVector) + states = first(trajectories) + data = Array{eltype(first(states))}( + undef, length(first(states)), length(states), length(trajectories)) + for (k, trajectory) in enumerate(trajectories) + for (t, state) in enumerate(trajectory) + data[:, t, k] = state + end end data end -data_tensor(states::AbstractVector) = data_tensor!(states_tensor(states), states, 1) - -# The tensor of one coordinate over the members of an ensemble, along the parameter axis. -function ensemble_data_tensor(solutions, coordinate::Symbol) - data = states_tensor(getproperty(first(solutions), coordinate), length(solutions)) - foreach(i -> data_tensor!(data, getproperty(solutions[i], coordinate), i), - eachindex(solutions)) - data -end - function data_tensors_from_geometric_solution( solution::GeometricSolution{ <:Number, <:Number, <:Any, <:NamedTuple{(:t, :q, :p, :q̇, :ṗ)}}) - (q = data_tensor(solution.q), p = data_tensor(solution.p)) + (q = data_tensor([solution.q]), p = data_tensor([solution.p])) end function data_tensors_from_geometric_solution(solution::GeometricSolution{ <:Number, <:Number, <:Any, <:Union{NamedTuple{(:t, :q, :v)}, NamedTuple{(:t, :q, :q̇)}}}) - data_tensor(solution.q) + data_tensor([solution.q]) end """ @@ -314,7 +303,7 @@ function DataLoader(ensemble_solution::EnsembleSolution{T, T1, Vector{ST}}; TT <: TimeSeries{T1}, ST <: Union{ GeometricSolution{T, T1, TT, NamedTuple{(:t, :q, :v), TuT}}, GeometricSolution{T, T1, TT, NamedTuple{(:t, :q, :q̇), TuT}}}} - DataLoader(ensemble_data_tensor(ensemble_solution.s, :q); + DataLoader(data_tensor([s.q for s in ensemble_solution.s]); autoencoder = autoencoder, suppress_info = suppress_info) end @@ -342,18 +331,18 @@ function DataLoader(ensemble_solution::EnsembleSolution{T, T1, Vector{ST}}; TT <: TimeSeries{T1}, ST <: GeometricSolution{T, T1, TT, NamedTuple{(:t, :q, :p, :q̇, :ṗ), TuT}}} DataLoader( - (q = ensemble_data_tensor(ensemble_solution.s, :q), - p = ensemble_data_tensor(ensemble_solution.s, :p)); + (q = data_tensor([s.q for s in ensemble_solution.s]), + p = data_tensor([s.p for s in ensemble_solution.s])); autoencoder = autoencoder, suppress_info = suppress_info) end -# `Base.copyto!` carries a host array to any backend; `KernelAbstractions.copyto!` has a method for -# the CPU alone, and a device backend such as `JLArray`'s has none. +# An array already on the backend is returned as it is. `Base.copyto!` carries a host array to any +# backend; `KernelAbstractions.copyto!` has a method for the CPU alone, and a device backend such as +# `JLArray`'s has none. function map_to_new_backend(input::AbstractArray{T}, backend::KernelAbstractions.Backend) where {T} - converted = KernelAbstractions.allocate(backend, T, size(input)...) - Base.copyto!(converted, input) - converted + networkbackend(input) == backend && return input + Base.copyto!(KernelAbstractions.allocate(backend, T, size(input)...), input) end # `map` over a `(q, p)` named tuple reaches each of the two tensors; over a plain array it would @@ -397,7 +386,7 @@ function DataLoader(dl::DataLoader{T1, <:QPTOAT, Nothing, K}, ) where {T1, K, T} kind = isnothing(autoencoder) ? K : data_set_kind(autoencoder) input = T == T1 ? dl.input : map_to_type(dl.input, T) - input = backend == networkbackend(dl) ? input : map_to_new_backend(input, backend) + input = map_to_new_backend(input, backend) DataLoader{T, typeof(input), Nothing, kind}( input, nothing, dl.input_dim, dl.input_time_steps, dl.n_params, nothing, nothing) diff --git a/src/data_loader/tensor_assign.jl b/src/data_loader/tensor_assign.jl index 5f6d9570a..435b471b2 100644 --- a/src/data_loader/tensor_assign.jl +++ b/src/data_loader/tensor_assign.jl @@ -1,7 +1,3 @@ -# Assigning a batch and its target out of the full data tensor is `src/data_loader/batch.jl`'s job: -# its gather kernel reads a matrix of `(time_step, parameter)` indices. This file crops the -# estimated output instead. - @doc raw""" assign_output_estimate(full_output, prediction_window) diff --git a/test/data_loader/data_loader.jl b/test/data_loader/data_loader.jl index 88c311512..5f6ca06bb 100644 --- a/test/data_loader/data_loader.jl +++ b/test/data_loader/data_loader.jl @@ -31,6 +31,15 @@ kind(dl::DataLoader) = typeof(dl).parameters[4] suppress_info = true))()) isa DataLoader @test (@inferred (() -> DataLoader((q = tensor, p = tensor); autoencoder = true, suppress_info = true))()) isa DataLoader + # ... and so do the matrix and vector constructors, which forward to them + vector = rand(T, 10) + @test (@inferred (() -> DataLoader(matrix; suppress_info = true))()) isa DataLoader + @test (@inferred (() -> DataLoader(vector; autoencoder = true, + suppress_info = true))()) isa DataLoader + @test (@inferred (() -> DataLoader((q = matrix, p = matrix); + suppress_info = true))()) isa DataLoader + @test (@inferred (() -> DataLoader((q = vector, p = vector); autoencoder = true, + suppress_info = true))()) isa DataLoader @test kind(DataLoader((q = tensor, p = tensor); suppress_info = true)) === :TimeSeries @test kind(DataLoader((q = rand(T, 10), p = rand(T, 10)); suppress_info = true)) === :TimeSeries diff --git a/test/quality/jet.jl b/test/quality/jet.jl index 5d0a79ec9..4c8edf506 100644 --- a/test/quality/jet.jl +++ b/test/quality/jet.jl @@ -125,7 +125,7 @@ const A3{T} = Array{T, 3} GML.convert_input_and_batch_indices_to_array, (TimeSeriesQP, Batch{:Transformer}, Indices); target_modules = TARGET))) # the one gather kernel of a `(q, p)` input: `gather` launches it once per tensor - @test isempty(kernel_body_reports(GML.assign_from_vector_of_tuples_kernel!, + @test isempty(kernel_body_reports(GML.gather_kernel!, (2, 3, 4), zeros(2, 3, 4), zeros(2, 8, 4), indices, 3)) # test/integration/docstrings/data_loader.jl @test isempty(JET.get_reports(JET.report_opt( @@ -150,7 +150,7 @@ const A3{T} = Array{T, 3} (DataLoader{T, A3{T}, A3{T}, :TimeSeries, Int}, BatchType, Indices); target_modules = TARGET))) # the gather kernel, once per element type - @test isempty(kernel_body_reports(GML.assign_from_vector_of_tuples_kernel!, + @test isempty(kernel_body_reports(GML.gather_kernel!, (2, 3, 4), zeros(T, 2, 3, 4), zeros(T, 2, 8, 4), indices, 3)) end # the `Batch` functor under `@allocated` in test/data_loader/batch_index_set.jl, and From 81bdfa8094585095c07dc121219ebbd050fba6c8 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 16 Oct 2026 15:44:40 +0200 Subject: [PATCH 8/8] Batch(batch_size) raises on a loader whose two time axes differ A Batch{:FeedForward} pairs the input and the output step by step. On a loader whose output has another time-axis length than its input, such as a classifier's with one output column, it gathered only the first patch of each sample and said nothing. number_of_start_positions now raises an ArgumentError that names Batch(batch_size, dl), and the batch test asserts it. The CHANGELOG migration sentence says so. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 +- src/data_loader/batch.jl | 10 +++++++++- test/data_loader/batch.jl | 4 +++- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d2e37fa4a..18d295bb5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -41,7 +41,7 @@ breaking release). ### Changed (breaking) -- **`DataLoader`'s `autoencoder` defaults to `false` in every constructor that takes it, except the copying constructor, where `autoencoder = nothing` inherits the loader's kind, and a loader given an input and an output is a set of samples.** A matrix or a vector is a time series unless `autoencoder = true` is passed: that keyword is what reads a matrix's second axis as the parameter index, and the matrix and the vector method defaulted to `true`, so migrate `DataLoader(matrix)` to `DataLoader(matrix; autoencoder = true)` where the columns are samples — `scripts/reproduction/symplectic_autoencoders/training.jl` and the tutorials do. `DataLoader(input, output)` builds a `:RegularData` loader whatever its kind, so its target is its output and never the input shifted by the sequence length: in 0.8.0 a `DataLoader(x, y)` of ten pairs drew the samples 1 to 9 and, where the input and the output differed in array type, returned one minibatch holding the whole data set. The count is `cld(n · n_params, batch_size)` with `n` the start positions of one parameter index, and a loader and a `Batch` that leave no start position raise an `ArgumentError` naming the time steps, the sequence length and the prediction window, where 0.8.0 asserted and an epoch could be empty. `Batch(batch_size, dl)` takes its two lengths from the loader, so a classifier's minibatch holds the whole patch sequence and the whole one-hot output of each of its samples. On a classifier's loader, `Batch(batch_size)` now draws only the first patch of each sample, where 0.8.0's minibatch held the whole input sequence, and nothing raises; migrate it to `Batch(batch_size, dl)`. `suppress_info` is a keyword of every constructor that takes a data set, which makes `DataLoader(solution, true)` a `MethodError`. +- **`DataLoader`'s `autoencoder` defaults to `false` in every constructor that takes it, except the copying constructor, where `autoencoder = nothing` inherits the loader's kind, and a loader given an input and an output is a set of samples.** A matrix or a vector is a time series unless `autoencoder = true` is passed: that keyword is what reads a matrix's second axis as the parameter index, and the matrix and the vector method defaulted to `true`, so migrate `DataLoader(matrix)` to `DataLoader(matrix; autoencoder = true)` where the columns are samples — `scripts/reproduction/symplectic_autoencoders/training.jl` and the tutorials do. `DataLoader(input, output)` builds a `:RegularData` loader whatever its kind, so its target is its output and never the input shifted by the sequence length: in 0.8.0 a `DataLoader(x, y)` of ten pairs drew the samples 1 to 9 and, where the input and the output differed in array type, returned one minibatch holding the whole data set. The count is `cld(n · n_params, batch_size)` with `n` the start positions of one parameter index, and a loader and a `Batch` that leave no start position raise an `ArgumentError` naming the time steps, the sequence length and the prediction window, where 0.8.0 asserted and an epoch could be empty. `Batch(batch_size, dl)` takes its two lengths from the loader, so a classifier's minibatch holds the whole patch sequence and the whole one-hot output of each of its samples. On a loader whose output's time axis differs from its input's, such as a classifier's with one output column, `Batch(batch_size)` raises an `ArgumentError` that names `Batch(batch_size, dl)`, where 0.8.0's minibatch held the whole input sequence; migrate it to `Batch(batch_size, dl)`. That batch is a `Batch{:Transformer}`, so `train!` on a `NeuralNetworkIntegrator` with it needs an explicit `loss`. `suppress_info` is a keyword of every constructor that takes a data set, which makes `DataLoader(solution, true)` a `MethodError`. - **The two output fields of `DataLoader` carry a fifth type parameter.** `output_dim` and `output_time_steps` are the type `S <: Union{Int, Nothing}` of the loader, so `DataLoader{T, AT, OT, :TimeSeries}` still names a loader and a *full* type name gains `, Nothing` for a loader without an output; a four-parameter constructor `DataLoader{T, AT, OT, K}(input, output, input_dim, input_time_steps, n_params, output_dim, output_time_steps)` fills `S` in from the fields. diff --git a/src/data_loader/batch.jl b/src/data_loader/batch.jl index 3647f03a5..db39962bd 100644 --- a/src/data_loader/batch.jl +++ b/src/data_loader/batch.jl @@ -153,11 +153,19 @@ number_of_batches(dl::DataLoader, batch::Batch) = cld( # The start positions of one parameter index: the input window must fit the input's time axis, and # the target must fit its own. The target is the output at offset `0` where the loader holds one, # the input at offset `seq_length` where it is a time series without an output, and there is none -# where the loader is a set of samples without an output. +# where the loader is a set of samples without an output. A `Batch(batch_size)` pairs the input and +# the output step by step, so it needs one length on the two time axes. function number_of_start_positions(dl::DataLoader{T, AT, OT, K}, batch::Batch) where {T, AT, OT, K} input_window = dl.input_time_steps - batch.seq_length + 1 target_window = if OT <: AbstractArray + if batch isa Batch{:FeedForward} && dl.input_time_steps != dl.output_time_steps + throw(ArgumentError( + "`Batch(batch_size)` pairs the input and the output time step by time step, " * + "and the `DataLoader` has $(dl.input_time_steps) time steps and " * + "$(dl.output_time_steps) in its output; `Batch(batch_size, dl)` draws the " * + "whole sequence of each sample.")) + end dl.output_time_steps - batch.prediction_window + 1 elseif K === :TimeSeries input_window - batch.prediction_window diff --git a/test/data_loader/batch.jl b/test/data_loader/batch.jl index 540a6f1ee..ad0aefae8 100644 --- a/test/data_loader/batch.jl +++ b/test/data_loader/batch.jl @@ -162,7 +162,9 @@ end @test batch == Batch(2, 4, 1) minibatches = batch(dl) @test length(minibatches) == 5 - @test length(Batch(2)(dl)) == 5 + # `Batch(2)` would pair the first patch with the one output column, so it raises + @test_throws ArgumentError Batch(2)(dl) + @test_throws "Batch(batch_size, dl)" number_of_batches(dl, Batch(2)) for indices in minibatches input_batch, output_batch = convert_input_and_batch_indices_to_array(dl, batch, indices) @test size(input_batch) == (9, 4, 2)