GML part M2: DataLoader samples against trajectories, one count and one gather

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` 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.
+
 - **`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 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.
+
 ### 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, 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
 
 - `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.

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`
@@ -376,6 +297,20 @@ 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: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))`. 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
 
 ### K3 · `test/quality/jet.jl` does not see an unstable kernel-launch argument, nor two barriers in the Poisson-tensor kernel

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 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

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