diff --git a/CHANGELOG.md b/CHANGELOG.md index 7dc1acf26..18d295bb5 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` 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. diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index eb936965a..bc059850d 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` @@ -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 diff --git a/docs/src/data_loader/data_loader.md b/docs/src/data_loader/data_loader.md index 22b34f1e6..86251f7c8 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 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/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..db39962bd 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. +# 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, :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,9 @@ 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 start positions of one +parameter index. No two indices are drawn twice, and the last minibatch holds what is left. # Examples @@ -148,99 +147,87 @@ 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)) -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)) -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)] +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. 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 + else + input_window end - batches[n_batches] = complete_indices[((n_batches - 1) * batch.batch_size + 1):end] - batches + check_window(min(input_window, target_window), dl, batch) 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) +# 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 && 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 -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) +# 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)))) + map(collect, Iterators.partition(complete_indices, batch.batch_size)) 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 gather_kernel!(gathered, data, indices, offset) i, j, k = @index(Global, NTuple) - input[i, j, k] = data[i, indices[1, k] + j - 1, indices[2, k]] + gathered[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}}) + map_to_new_backend([index[i] for i in 1:2, index in batch_indices_tuple], 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. +function gather(data::AbstractArray{T, 3}, indices::AbstractMatrix{Int}, offset::Int, + width::Int) where {T} + backend = networkbackend(data) + gathered = KernelAbstractions.allocate( + backend, T, size(data, 1), width, size(indices, 2)) + gather_kernel!(backend)(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}}, indices::AbstractMatrix{Int}, + offset::Int, width::Int) where {AT <: AbstractArray{<:Any, 3}} + map(tensor -> gather(tensor, indices, offset, width), data) end @doc raw""" @@ -253,7 +240,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 +260,18 @@ 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} + indices = convert_vector_of_tuples_to_matrix(networkbackend(dl), batch_indices_tuple) + 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 OT <: AbstractArray + (input, gather(dl.output, indices, 0, batch.prediction_window)) + elseif K === :TimeSeries + (input, gather(dl.input, indices, batch.seq_length, batch.prediction_window)) + else + input + end end diff --git a/src/data_loader/data_loader.jl b/src/data_loader/data_loader.jl index b4827fb06..f5f402434 100644 --- a/src/data_loader/data_loader.jl +++ b/src/data_loader/data_loader.jl @@ -29,32 +29,41 @@ 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, +# 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 +# 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 +data_set_shape(autoencoder::Bool, n::Int) = autoencoder ? (1, n) : (n, 1) + """ DataLoader(data::AbstractArray{<:Number, 3}) @@ -90,7 +99,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 +116,19 @@ 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( +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." 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,20 @@ 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} +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) 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) + @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) + DataLoader(reshape(data, input_dim, time_steps, n_params); + autoencoder = autoencoder, suppress_info = true) end """ @@ -169,20 +168,25 @@ 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 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 <: Union{T, Integer}} @assert size(input, 3) == size(output, 3) if !suppress_info @info "You have provided an input and an output." end - 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 @@ -209,83 +213,62 @@ 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 @info "You have provided a NamedTuple with keys q and p; the data are matrices. This is interpreted as *symplectic data*." end - - 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 + dim2, second_dim = size(data.q) + 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}}; +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 @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 -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}} - 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] +# 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 -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}}} - 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 +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 - data +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 """ @@ -306,45 +289,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(data_tensor([s.q for s in ensemble_solution.s]); + autoencoder = autoencoder, suppress_info = suppress_info) end """ @@ -367,52 +327,33 @@ 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 = 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 -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₂ +# 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} + networkbackend(input) == backend && return input + Base.copyto!(KernelAbstractions.allocate(backend, T, size(input)...), input) 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,42 +379,20 @@ 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; + ::Type{T} = 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, T} + kind = isnothing(autoencoder) ? K : data_set_kind(autoencoder) + input = T == T1 ? dl.input : map_to_type(dl.input, T) + 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...) +function DataLoader(dl::DataLoader, ::Type{T}; kwargs...) where {T} DataLoader(dl, networkbackend(dl), T; kwargs...) end diff --git a/src/data_loader/tensor_assign.jl b/src/data_loader/tensor_assign.jl index f351bc9ec..435b471b2 100644 --- a/src/data_loader/tensor_assign.jl +++ b/src/data_loader/tensor_assign.jl @@ -1,22 +1,3 @@ -""" -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 - @doc raw""" assign_output_estimate(full_output, prediction_window) @@ -31,52 +12,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..8da9c4f7b 100644 --- a/test/data_loader/accuracy.jl +++ b/test/data_loader/accuracy.jl @@ -1,18 +1,15 @@ +# `accuracy` is documented API -- `docs/src/data_loader/data_loader.md` renders both its methods -- +# 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)`. 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 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. +using Zygote Random.seed!(1234) @@ -29,9 +26,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 +53,23 @@ end GeometricMachineLearning.accuracy(model, ps, dl) end end + +# 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..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 @@ -25,6 +20,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 @@ -32,16 +30,13 @@ 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) 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..ad0aefae8 --- /dev/null +++ b/test/data_loader/batch.jl @@ -0,0 +1,204 @@ +# 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. + +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 + @test number_of_batches(dl, batch) == 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 eltype(output) === 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) + 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 (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 + +# 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) + @test length(minibatches) == 1 + @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 + +# 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 + @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) + + # 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 + +# 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 + # `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) + @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 + +# 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 ($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{T, 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) + 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/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..5f6ca06bb --- /dev/null +++ b/test/data_loader/data_loader.jl @@ -0,0 +1,208 @@ +# The constructors of `src/data_loader/data_loader.jl`: the `autoencoder` default and the keyword +# 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 +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 + # 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 + # ... 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 + # ... 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 + +# 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( + 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 + + # 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!` +# 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 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) + 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 + +# 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) + @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..14a571375 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 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..4c8edf506 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.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( 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.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 @@ -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")