Skip to content

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

Merged
michakraus merged 8 commits into
mainfrom
simplify/part-m2
Oct 16, 2026
Merged

michakraus merged 8 commits into
mainfrom
simplify/part-m2

Conversation

@michakraus

@michakraus michakraus commented Oct 16, 2026 •

Copy link
Copy Markdown
Member

The DataLoader half of the GeometricMachineLearning 0.9 simplification — part M2.
A supervised loader becomes a set of samples, the two batching defects go, every constructor that
takes a data set defaults autoencoder = false, and the batching code is a little over half its
former size.

Fixes #211
Relates to #212: the batch indices reach the backend in one copy, and the @views it asked for
are declined (see One gather below).
Refs #325 (the batch half of the JET dispatch reports; the PoissonTensor half is part M6's and
this change uses no closing keyword for it)

What this changes

  • A supervised DataLoader(input, output) holds independent samples, whatever its kind. At the
    base it was typed :TimeSeries and drew the samples 1 to 9 of ten pairs, and where the input and
    the output differed in array type one minibatch held the whole data set. It is now :RegularData
    and draws every pair once per epoch. Batching branches on the output type first: an array output
    gives the pair (input[:, t:t+sl-1, p], output[:, t:t+pw-1, p]) at offset 0, no output with
    :TimeSeries gives the input at offset seq_length, and no output with :RegularData gives the
    input window alone.
  • One count. Per parameter index the start positions are 1:n with n = min(input_time_steps − seq_length + 1, target_steps − offset − prediction_window + 1), and the batch count is
    cld(n · n_params, batch_size). n < 1 raises an ArgumentError naming the time steps,
    seq_length and prediction_window, where the base asserted. The draw order is unchanged, so a
    seeded trajectory run draws the minibatches it drew before, except the ones the last sample joins.
  • Batch(batch_size) needs one time-axis length. A Batch{:FeedForward} pairs the input and the
    output step by step, so on a loader whose output's time axis differs from its input's — a
    classifier's, with one output column — it raises an ArgumentError that names
    Batch(batch_size, dl), where it would train each sample on its first patch. Batch(batch_size, dl) is a Batch{:Transformer}, so train! on a NeuralNetworkIntegrator with it needs an
    explicit loss.
  • One gather. The six convert_input_and_batch_indices_to_array methods and four kernels become
    one @kernel, out[i, j, k] = data[i, idx[1, k] + offset + j − 1, idx[2, k]], and one helper that
    applies it to an array or to a q/p pair; the types of out and data are independent, because
    a view output is a ReshapedArray of a SubArray. The index matrix is one host-built 2 × n
    Matrix{Int} copied to the backend once, and there are no @views — a list of (t, p) windows is
    not a strided view, which is what Don't allocate arrays in convert_input_and_batch_indices_to_array but use @views. #212 had asked for.
  • One autoencoder default. It is false in every constructor that takes a data set, with the
    matrix and vector methods' true gone, and suppress_info is a keyword of every one of those
    constructors. The copying constructor DataLoader(dl, backend, T) keeps autoencoder = nothing,
    meaning inherit. The manual, the transformer tutorial and the reproduction script that read a
    matrix as a set of samples now pass autoencoder = true explicitly.
  • Possible type instability in definition of DataLoader #211. output_dim and output_time_steps carry a fifth type parameter S <: Union{Int, Nothing}; K stays the fourth, so DataLoader{T, AT, OT, :TimeSeries} patterns still match. A
    four-parameter constructor DataLoader{T, AT, OT, K}(…) fills S in from the fields, and a kernel
    launch no longer reads a value of no single type.
  • C17. DataLoader(input::AbstractArray{T, 3}, output::AbstractArray{T1, 3}) where {T, T1 <: Union{T, Integer}} is one method, so the output's element type is the input's or an integer and
    accuracy is reachable through a public constructor. detect_ambiguities finds none, an Int/Int
    pair builds :RegularData, and a Float32 input with a Float64 output stays a MethodError.
  • tensor_assign.jl. assign_output_estimate(x, pw) = x[:, (end − pw + 1):end, :] with no rrule:
    Zygote's getindex 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 kernel, augment_zeros and the file-top string Julia discards go.
  • B10 stays, and is tested. The DataLoader(::EnsembleSolution{T, T1, Vector{ST}}) method for a
    dataser of (t, q, q̇) is reachable — a rigid-body odeensemble produces one, and the tutorial
    and three scripts call it — so it is not deleted; test/data_loader/data_loader.jl tests it.
  • JET reports a runtime dispatch in PoissonTensor and in one batch method #325's batch half. 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 size of the batching code

Counted as code lines — a code line holds a token that is not whitespace, not a comment
and not inside a docstring — over the three files, before and on the branch:

file before on the branch
batch.jl 196 92
data_loader.jl 256 174
tensor_assign.jl 49 4
total 501 270

270 is 53.9 % of 501, inside the part's ceiling of 275.

Tests

The figures in this section were measured before the review fixes; the review thread states what
was run on those.

The part's own mutants, run warm against the test files that must catch them: 34 mutants, 34
CAUGHT, 0 SURVIVED, 0 UNKNOWN, 0 INVALID
(and the 26 of the build round, 26 CAUGHT). The named
units pass: data_loader/batch.jl, data_loader/data_loader.jl, data_loader/accuracy.jl,
data_loader/autoencoder_keyword.jl, data_loader/batch_index_set.jl,
data_loader/batch_data_loader_qp_test.jl, data_loader/batch_functor.jl,
data_loader/optimizer_functor_with_adam.jl, integration/docstrings/data_loader.jl,
kernels/kernel_pullbacks.jl and the doctests group — 1970 assertions, 0 failures.

The four whole-package runs, all green at a5e959e0, whose diff to the critic's head 07ab13e9 is
comment lines only with an identical code and docstring AST:

run result
full suite, Julia 1.13.2 PASS, exit 0, 594.9 s — Testing GeometricMachineLearning tests passed, the only non-pass being JET's one known broken PoissonTensor test
full suite, the 1.12.7 floor PASS, exit 0, 871.2 s — the same totals line
docs build (julia --project=docs docs/make.jl) PASS, exit 0, 1160 s — 0 unresolved references
Scripts smoke run of the symplectic-autoencoder training.jl with GML_SMOKE=1 PASS, exit 0, 80 s

The head carries, after those runs, the pre-PR verifier's fixes (an inference fix on the tensor
constructors, five comment corrections, the KNOWN_ISSUES.md K18 entry) and the correction of the
CHANGELOG entry's domain. The full suite at the head is CI's, and the same command is left to the
dispatcher of the part.

Critic

Final critic round on the part (a critic council member of the build loop), judged a5e959e0..07ab13e9:

Verdict: PASS.

clause domain on origin/main on the branch verdict
Exact seven dictated lines accuracy.jl header required header absent exact-header assertion 1/1 PASS
(a) precision comment contract header describing constructor types 2/2 2/2 PASS
(b) prohibited phrases absent; replacement once src, test, accuracy header called it. remains in the baseline accuracy comment; replacement count 0 no prohibited matches; replacement once at line 6 PASS
(c) only permitted file changes previous-head diff not applicable only test/data_loader/accuracy.jl; 1/1 PASS
(d) only comments change diff lines and Julia AST, excluding source locations not applicable five added and five removed comment lines; identical code/docstring AST; 3/3 PASS
(e) added lines ≤100 characters added diff lines not applicable maximum 94 characters; 1/1 PASS
(f) accuracy unit passes 16/16 data_loader/accuracy.jl 11 pass, 1 fail, 3 error 16/16 PASS
Dictated precision description matches behavior Float32/Float64 inputs; same-type, Int and Bool targets; mismatched float types not rerun: unchanged code 16/16 PASS
Header caller claim matches source package, extension and script references not rerun: unchanged source nine references in one source file, inspected; the live caller is accuracy's NeuralNetwork forwarding method PASS
Reached quality checks six selected files not applicable 126 pass, 1 known broken PASS
Earlier evidence gaps remain closed former promotion, reversed-column and coordinate-swap survivors not rerun: previously closed defects 9 VALID; smoke CAUGHT; 9 CAUGHT, 0 SURVIVED, 0 UNKNOWN, 0 INVALID PASS
Earlier backend, allocation, history and defaults defects remain closed earlier reproducer inputs not rerun: previously closed defects as below PASS

Per-file test counts of that round:

file passes broken
data_loader/accuracy.jl 16 0
quality/reachability.jl 2 0
quality/exports.jl 8 0
quality/aqua.jl 17 0
quality/jet.jl 77 1
quality/docs_make_checks_references.jl 6 0
Total 126 1

Every defect of the earlier rounds, as the same round reported it:

round defect now
1a 1a-1: JLArray batching fails before gather fixed: 9/9 per precision under allowscalar(false)
1a 1a-2: output and p precision insufficiently tested fixed: all CAUGHT
1a 1a-3: (q, p) gathered columns unchecked fixed: CAUGHT
1a 1a-4: single-solution coordinate swap undetected fixed: CAUGHT
1a 1a-5: historical comments and false precision claim fixed: history probe 3/3; no named historical phrases; precision probe 2/2; exact dictated header
1a 1a-6: manual misdescribes second example fixed: 11/11
1c 1c-1: Float32 target promotion survives fixed: CAUGHT
1c 1c-2: redundant ensemble-coordinate allocations fixed: actual 3/3, excess 576 B; synthetic 2/2, excess 0 B in both precisions
1c 1c-3: implementation history in comments fixed: 3/3
2 remaining false "two element types are free" claim fixed: 2/2; independent runtime contract 16/16
2, assigned other defect introductory called it. history fixed: prescribed present-tense text; no prohibited match

Blocking defects: None. Other defects: None. Fix first: none.

Pre-PR verification

Verdict: clear to open — 0 unresolved.

Verified simplify/part-m2 at 07ab13e9, including its pending working-tree changes, against merge
base d751ecd9.

  • Fixed: tensor-constructor inference is restored with Base.@constprop :aggressive alone,
    without duplicated kind branches. Earlier comment and K18 corrections remain intact. The changelog
    now states the copying-constructor inheritance exception and correctly limits the suppress_info
    claim.
  • Unresolved: none.
  • Pre-existing: the manual sampling formula is one short; nearby prose says "with" instead of
    "without" replacement; two docstrings misdescribe index types or unconditional allocation. These
    remain recorded and do not block this change.
  • Checked and clean: cold inference probes 8/8, constructor tests 724/724, keyword tests 7/7,
    batching tests 300/300, and index-allocation regressions 2/2 on both Julia 1.13.2 and the 1.12.7
    floor. Live code_typed checks return concrete constructor and CPU minibatch types. Aqua piracy
    passes. Static lint reports 16 changed Julia files clean. Annotation-expansion checks confirm
    unchanged executable ASTs, and changelog contract assertions pass 14/14. The three-file code-line
    count is 90 + 180 + 4 = 274, within the limit of 275.
  • Not checked: full suites, docs build and smoke script were not repeated. No new whole-package
    JET/import sweep, solution/ensemble-specific inference probe, GPU measurement or timing benchmark
    was performed.

Known issues added

KNOWN_ISSUES.md gains one entry, K18 — the tensor-sampling formula in the DataLoader manual
is one short: item 1 of Sampling from a tensor gives the start positions as
shuffle(1:(input_time_steps − seq_length − prediction_window)), where the code takes
n = input_time_steps − seq_length − prediction_window + 1 and counts cld(n · n_params, batch_size) batches (three start positions against the code's four for the page's five time steps
and Batch(4, 2, 0)). It is a docs page the part's own text pass does not own, so it is recorded
rather than widened into this change.

The branch also removes four entries, because this change closes them: B10 (the ensemble
method has a caller and now a test, so it stays), C17 (accuracy is reachable through the
public two-element-type constructor), K8 (a loader of one sample batches) and K12 (Batch
now draws the last sample of a loader of matrices).

Reported, not fixed

Pre-existing defects the pre-PR verifier found beside the diff. They are outside the part's
sections, so they are recorded here and not fixed in this change:

file:line claim
docs/src/data_loader/data_loader.md:143 K18 above: the tensor-sampling upper bound needs + 1. The base has the same formula; recorded as K18.
docs/src/data_loader/data_loader.md:98 "With replacement" should read "without replacement": the base file has the same sentence, and the batching partitions the shuffled Cartesian product without repeating indices.
src/data_loader/batch.jl:12 The docstring says vectors of integers rather than vectors of index tuples; a code-typed probe returns Vector{Vector{Tuple{Int64,Int64}}}. The sentence is outside the changed lines.
src/data_loader/data_loader.jl:363 The copying-constructor docstring promises allocation unconditionally, where a matching backend and element type retain the original tensor; the base carries the same identity branches.

…art M2)

A loader built from an input and an output holds a set of samples whatever its
kind, so batching no longer branches on the kind alone: the target is the output
itself, never the input shifted by the sequence length. One count, nn = the start
positions of one parameter index, replaces the two the batching had, and it draws
every sample of a supervised loader including the last. The five gather methods
and four kernels become one kernel and one helper; the copy loops of a solution
and an ensemble become one `data_tensor`; `map_to_*` reaches the two tensors of a
`(q, p)` pair with `map`; `tensor_assign.jl` keeps only the crop of the output.

autoencoder defaults to false in every constructor, and `DataLoader(input,
output)` takes two free element types, so a real input with an integer one-hot
target is one call. The two output fields of the struct carry a fifth type
parameter, which makes them concrete. KNOWN_ISSUES B10, C17, K8 and K12 leave the
file: the ensemble method is reachable through a rigid-body `odeensemble`, and it
has a test.
The autoencoder default, the supervised kind, the fifth type parameter and the
new element-type constructor; B10, C17, K8, K12 and the batch half of #325.
The advisor ruling of 2026-10-16: the index set through Iterators.partition, the
kernel signature without its type parameters, the index matrix in one
comprehension, one if/elseif/else in the launcher, one data_set_shape, one
(input, output) method with T1 <: Union{T, Integer}, the copy methods dispatching
on the dataser keys, the matrix and (q, p) methods forwarding to the tensor
method, if !suppress_info at the five sites, and ::Type{T} at the last two.
…unts

- `map_to_new_backend` copies with `Base.copyto!`, so the host index matrix reaches a JLArray:
  `KernelAbstractions.copyto!` has a method for the CPU alone.
- the four-parameter constructor converts its five counts, so the `Int32` sizes `GMLDatasets`
  passes build a loader.
- `check_window` names the input's and the output's time steps and the parameter count, and it
  rejects an empty parameter axis, which drew an empty epoch in silence.
- `ensemble_data_tensor` fills the final tensor through one coordinate copy, so no whole
  coordinate tensor is built per member.
- the tests assert the element type of a supervised minibatch and of a `(q, p)` target, compare
  the gathered columns of a `(q, p)` loader and of one `(q, p)` solution with their own data,
  and run the one-sample and typed-kind rows in both precisions.
- the comments state the current code, and the manual states the `autoencoder` default correctly.
- `KNOWN_ISSUES.md` gains K18, the tensor-sampling formula that is one short.
The two tensor constructors keep data_set_kind(autoencoder) and gain Base.@constprop :aggressive, which restores the concrete kind the base inferred, without the duplicated branches that took the three-file code-line count to 284. The count is 274 of the part's 275. Six comments now state the current contracts, KNOWN_ISSUES.md gains K18 for the manual's sampling formula, and the CHANGELOG entry no longer claims the autoencoder default and suppress_info for the copying constructor.
@codecov

codecov Bot commented Oct 16, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 95.23810% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.29%. Comparing base (d751ecd) to head (81bdfa8).

Files with missing lines Patch % Lines
src/data_loader/batch.jl 94.44% 2 Missing ⚠️
src/data_loader/data_loader.jl 95.65% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #346      +/-   ##
==========================================
+ Coverage   88.56%   91.29%   +2.73%     
==========================================
  Files          65       65              
  Lines        2212     2080     -132     
==========================================
- Hits         1959     1899      -60     
+ Misses        253      181      -72     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@michakraus michakraus left a comment

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

JuliaGNI/GeometricMachineLearning.jl PR #346 — GML part M2: DataLoader samples against trajectories, one count and one gather

Verdict: request changes (posted as a comment review: GitHub refuses a review state on one's own PR)
Reviewed: 6282628
CI: every job green except Downgrade - ubuntu-latest, which is a resolver failure and fails on main too. It dies in julia-downgrade-compat before any test runs:

Error: forcedeps check failed: GPUArraysCore resolved to 0.2.1 but lower bound is 0.2.0
Error: forcedeps check failed: Zygote resolved to 0.7.5 but lower bound is 0.7.0
Error: forcedeps check failed: ChainRulesCore resolved to 1.25.1 but lower bound is 1.0.0
Error: forcedeps check failed: NNlib resolved to 0.9.30 but lower bound is 0.8.0
Error: forcedeps check failed: KernelAbstractions resolved to 0.9.38 but lower bound is 0.9.0
ERROR: LoadError: forcedeps check failed for .: Some packages did not resolve to their lower bounds.

The same five lines are in the Downgrade job of main at d751ecd (run 38854142587, job 116821714018), whose workflow still concludes success. The PR does not touch Project.toml.

Findings

# severity file:line claim evidence fix inside a changed hunk?
1 correctness-risk src/data_loader/batch.jl:157-168 On a classifier loader (input 4 time steps, output 1), Batch(batch_size) now gives n = min(4, 1) = 1 start position, so each sample trains on its first patch only. The base gathered the whole sequence. No error is raised, and the CHANGELOG gives no migration line. The PR's own test had to change to Batch(2, dl) (test/data_loader/optimizer_functor_with_adam.jl:41). PR: DataLoader(rand(Float32, 9, 4, 10), zeros(Int, 10, 1, 10)), Batch(2) → input size (9, 1, 2), == input[:, 1:1, …]: true. Base, same data through the parametric :TimeSeries constructor → input size (9, 4, 2). yes
2 quality src/data_loader/data_loader.jl:150, :174, :216, :240 The matrix, vector, (q, p)-matrix and (q, p)-vector constructors infer DataLoader{Float32, Array{Float32, 3}, Nothing, _A, Nothing} where _A, even with literal keywords. The base inferred a small Union of the two concrete kinds. @constprop is only on the two tensor methods these forward to, and the new @inferred rows (test/data_loader/data_loader.jl:27-33) test only those two. Base.infer_return_type(x -> DataLoader(x; suppress_info = true), (Matrix{Float32},)) → the where _A type on the PR. The base gives Union{DataLoader{…, :RegularData}, DataLoader{…, :TimeSeries}}. With Base.@constprop :aggressive on the matrix method, both the matrix and the vector call infer …, :TimeSeries, Nothing}. yes
3 quality src/data_loader/data_loader.jl:247-271 Four helpers (states_tensor, data_tensor!, data_tensor, ensemble_data_tensor) and a foreach closure do one copy loop. ensemble_data_tensor(solutions, coordinate::Symbol) infers Array{_A, 3} where _A, so both ensemble constructors return a non-concrete type. The base was also non-concrete there (Any), so this is not a regression, but the rewrite keeps the defect. infer_return_type of ensemble_data_tensor(e.s, :q) → Array{_A, 3} where _A; DataLoader(ensemble) → DataLoader{T, AT, Nothing, _B, Nothing} where …. The form below gives equal tensors in 3 cases (a (q, q̇) ensemble, a (q, p) solution, a (q, p) ensemble), infers Array{Float64, 3}, and allocates 320 B where the current code allocates 560 B on the 2-member rigid-body ensemble. yes
4 quality src/data_loader/batch.jl:208-225, :262-271 gather takes dim and batch_size, but the arrays already hold both values: size(data, 1) and size(indices, 2). Removing them also removes the ÷ 2, the T/AT parameters of the NamedTuple method, the batch_size local of the launcher, and the last sentence of the comment at :208-211. The form below is == to gather for a tensor, a (q, p) pair and a JLArray under allowscalar(false). input_dim == size(input, 1) and output_dim == size(output, 1) hold for the tensor, matrix, (q, p) and (input, output) constructors. yes
5 nit src/data_loader/batch.jl:195 The name assign_from_vector_of_tuples_kernel! describes an input that the kernel does not take: per the comment at :201, it takes a 2 × n matrix. read yes (rename to gather_kernel!, also at test/quality/jet.jl:128, :153)
6 nit src/data_loader/batch.jl:205, src/data_loader/data_loader.jl:352-357, :400 On the CPU, each minibatch builds its host index matrix and then copies it into a second host matrix. The copying constructor holds the identity branch that would prevent this. @allocated for 256 indices: helper 8352 B, bare comprehension 4176 B yes
7 nit src/data_loader/batch.jl:160 vs :265; :266-271 "Has an output" is tested in two forms (OT <: AbstractArray, dl.output !== nothing). Also, a tuple return without parentheses runs over two lines (input, / gather(…)), which reads as two statements. read yes
8 nit src/data_loader/data_loader.jl:184, test/data_loader/data_loader.jl:197 This validates the caller's arguments with @assert, and a new test pins the AssertionError, in a change that replaces the batch @assert with an ArgumentError. read yes
9 nit src/data_loader/data_loader.jl:194-202 The matrix and vector (input, output) methods still require one T. So a Float32 input with an Int target is a MethodError for these shapes, but CHANGELOG:97 says that DataLoader(input, output) "accepts two element types". DataLoader(rand(Float32, 4, 6), zeros(Int, 3, 6)) → MethodError; vector likewise; the tensor form → DataLoader{Float32, …, Array{Int64, 3}, :RegularData, Int64} no: the lines are unchanged
10 nit src/data_loader/tensor_assign.jl:1-3 The file-top comment describes batch.jl, not this file. read yes
11 nit CHANGELOG.md:44, :80 A CHANGELOG reader cannot resolve "at the base" (:44) or "in the two files" (:80). read yes
12 nit PR body Fixes #212 closes an issue that asks for @views (closingIssuesReferences = 211, 212), and the body declines them. "Advisor decisions" has two headings with no text under them, and it cites a plan that is not in the PR. gh pr view 346 --json closingIssuesReferences; gh issue view 212 —

Concrete fixes

1. Make the case loud. In number_of_start_positions, when the loader has an output whose time axis differs from the input's and the batch is Batch{:FeedForward}, throw an ArgumentError that names Batch(batch_size, dl). Pin it with a @test_throws beside test/data_loader/batch.jl:165, which now asserts length(Batch(2)(dl)) == 5. If you keep the current behaviour, add a migration sentence to CHANGELOG.md:44: Batch(batch_size) on such a loader draws a window of length one at offset 0, so use Batch(batch_size, dl) for a classifier.

2. Add Base.@constprop :aggressive to the matrix method (:150) and to the (q, p)-matrix method (:216). The vector methods forward to these two. Then add @inferred rows for DataLoader(matrix) and DataLoader((q = matrix, p = matrix)).

3. One function over a vector of trajectories:

# The tensor of one coordinate over a vector of trajectories, with the axes `(system dimension,
# time steps, trajectories)`. A trajectory's time axis is `0:n`, so the states are copied 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 (i, trajectory) in enumerate(trajectories), (t, state) in enumerate(trajectory)
        data[:, t, i] = state
    end
    data
end

The callers become data_tensor([solution.q]), (q = data_tensor([solution.q]), p = data_tensor([solution.p])), data_tensor([s.q for s in ensemble_solution.s]) and its p twin.

4.

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

gather(data::QPT, args...) = map(tensor -> gather(tensor, args...), data)

The call sites become gather(dl.input, indices, 0, batch.seq_length) and so on.

6.

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

Then :400 is input = map_to_new_backend(input, backend). Base.copyto! returns its destination on Array and on JLArray (probed).

7. Branch on OT <: AbstractArray in both places, and return (input, gather(…)).

8. size(input, 3) == size(output, 3) || throw(DimensionMismatch("…")), and @test_throws DimensionMismatch.

9. where {T, T1 <: Union{T, Integer}} with output::AbstractMatrix{T1} and output::AbstractVector{T1}.

10. Delete the three lines.

11. Change "at the base" to "in 0.8", and delete "in the two files".

Verified and fine

  • detect_ambiguities(GeometricMachineLearning; recursive = false): 0.
  • test/data_loader/batch.jl and test/data_loader/data_loader.jl, included in a session on Julia 1.13.2: 1024 pass, 0 fail.
  • The tensor and (q, p)-tensor constructors and a (q, p) GeometricSolution infer concrete types, with the kind fixed by a literal keyword.
  • The comment at data_loader.jl:350-351 holds: KernelAbstractions defines copyto! for CPU alone (src/cpu.jl:19), and JLArrays defines only Base.copyto!.
  • The deleted rrule: test/data_loader/accuracy.jl:67-75 checks Zygote's adjoint, zero on the cropped columns and with the primal's element type.
  • JuliaFormatter (sciml) passes on the 7 changed src/test files, and test-layout.jl --check is clean.
  • No text in the PR is addressed to a reviewing agent. The body's Critic and Pre-PR verification sections are the author's own reports, and this review does not use them as evidence.

Not checked

  • The full suite, the docs build and Metal. CI ran them, and they are green.
  • The cold timing of the gather against the base. The load average was 12.8, with 32 Julia processes running. The (q, p) launcher now launches 4 kernels where the base launched 2.
  • The GMLDatasets caller that data_loader.jl:51-53 and the tests name. It is not in this tree.
  • The PR body's 274-code-line count and its mutation-test figures.

…te matrix loaders

- gather takes its sizes from the arrays it reads: the first axis from
  the data, the batch size from the index matrix. The dim, batch_size and
  the halving for a (q, p) pair go.
- The gather kernel is gather_kernel!, and the batch functor tests the
  output by one form, OT <: AbstractArray, as number_of_start_positions does.
- data_tensor builds the tensor of one coordinate over a vector of
  trajectories; states_tensor, data_tensor! and ensemble_data_tensor go.
- map_to_new_backend returns an array already on the backend, so a CPU
  minibatch no longer copies its index matrix, and the copying
  constructor drops its own backend branch.
- The matrix and (q, p)-matrix constructors carry
  Base.@constprop :aggressive, so the matrix and vector constructors infer
  one concrete loader for a literal autoencoder; the tests assert it.
- tensor_assign.jl loses the file-top comment that described batch.jl.
- CHANGELOG: "the base" is 0.8.0, Batch(batch_size) on a classifier's
  loader gets its migration, and the two-element-type claim names the
  3-tensor method.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
@michakraus

Copy link
Copy Markdown
Member Author

Follow-up to the review above. The fixes are in 6342d90. Every finding by its number:

# severity action detail
1 correctness-risk needs a decision, CHANGELOG migration added Batch(batch_size) on a classifier's loader still gathers input[:, 1:1, p], the first patch only (reproduced: input (9, 1, 2)). Raising an ArgumentError for a Batch{:FeedForward} whose loader has an output with a different time-axis length changes which loaders are valid, so the code is left to the author. The CHANGELOG entry for the autoencoder default now carries the migration to Batch(batch_size, dl). Note that Batch(batch_size, dl) is a Batch{:Transformer}, so train! on a NeuralNetworkIntegrator with it needs an explicit loss.
2 quality fixed Base.@constprop :aggressive on the matrix and the (q, p)-matrix constructors. DataLoader(matrix), DataLoader(vector; autoencoder = true) and their (q, p) forms infer one concrete loader type; test/data_loader/data_loader.jl asserts all four with @inferred.
3 quality fixed One data_tensor(trajectories) replaces states_tensor, data_tensor!, data_tensor and ensemble_data_tensor. The callers pass [solution.q] and [s.q for s in ensemble_solution.s]. The tensors are checked against solution.q[t] in test/data_loader/data_loader.jl, which passes.
4 quality fixed gather(data, indices, offset, width) takes the first axis from the data and the batch size from size(indices, 2); the (q, p) method maps over the two tensors with no ÷ 2. 25 minibatches over five loaders and five Batches are identical before and after.
5 nit fixed The kernel is gather_kernel!, also in test/quality/jet.jl.
6 nit fixed map_to_new_backend returns an array already on the backend, so the CPU index matrix is not copied, and the copying constructor drops its own backend branch. The byte figures of the review were not re-measured and are not claimed here.
7 nit fixed The batch functor tests OT <: AbstractArray, as number_of_start_positions does, and returns parenthesised tuples.
8 nit not fixed The @assert line of DataLoader(input, output) is unchanged by this PR, so it is outside the fix scope. One-line fix: size(input, 3) == size(output, 3) || throw(DimensionMismatch(…)), and @test_throws DimensionMismatch in the test.
9 nit not fixed in code, CHANGELOG corrected The matrix and vector (input, output) methods are unchanged lines. The CHANGELOG entry now says the two-element-type method is the 3-tensor one, and that the matrix and vector forms take one T.
10 nit fixed The file-top comment of tensor_assign.jl is gone.
11 nit fixed The CHANGELOG says 0.8.0 where it said "the base" (the data-loader sources are the same at v0.8.0 and the merge base d751ecd9), and "in the two files" is gone.
12 nit needs a decision The PR body is the author's: Fixes #212 closes #212, which asked for @views that the PR declines, and the "Advisor decisions" section has two empty headings. Left unchanged.

Checks on 6342d90: JuliaFormatter and fatou lint clean on the changed files; the pre-commit hook (layout, format, lint, NFC, load) passed. 19 test files — every file in test/data_loader/, quality/jet.jl, integration/docstrings/data_loader.jl, reduced_system/reduced_system.jl, integration/reduced_loss_optimization.jl, architectures/symplectic_autoencoder_tests.jl and architectures/psd_architecture_tests.jl — pass on Julia 1.13.2: 1367 pass, 0 fail, 0 error, 1 broken (the known PoissonTensor JET line, #325). The full suite, the doctests and the docs build are CI's.

CI before this push: the one red job, Downgrade - ubuntu-latest, fails in the resolver (forcedeps check failed: … did not resolve to their lower bounds) with no test run; the same lines are in that job on main at d751ecd9, and Project.toml is not in the diff. It is not a required check.

A Batch{:FeedForward} pairs the input and the output step by step. On a
loader whose output has another time-axis length than its input, such as
a classifier's with one output column, it gathered only the first patch
of each sample and said nothing. number_of_start_positions now raises an
ArgumentError that names Batch(batch_size, dl), and the batch test asserts
it. The CHANGELOG migration sentence says so.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
@michakraus

Copy link
Copy Markdown
Member Author

Update on the two findings that were open in the follow-up above, both decided by the author:

# action detail
1 fixed in 81bdfa8 number_of_start_positions raises an ArgumentError that names Batch(batch_size, dl) when a Batch{:FeedForward} meets a loader whose output's time axis differs from its input's. test/data_loader/batch.jl asserts it on the classifier loader, where it previously asserted the silent count. The CHANGELOG migration sentence says the call raises. A loader with one length on both time axes, and every Batch{:Transformer}, are unaffected.
12 fixed in the PR body #212 is now "relates to" (closingIssuesReferences is [211]), the empty "Advisor decisions" section is gone, and the code-line table gives the head's counts: 92 + 174 + 4 = 270 of 501.

Checks on 81bdfa8: the pre-commit hook (layout, format, lint, NFC, load) passed. 28 test files — every file in test/data_loader/, quality/jet.jl, integration/docstrings/data_loader.jl, integration/losses_and_optimization.jl, integration/reduced_loss_optimization.jl, integration/train_manual_architectures.jl, integration/mixed_precision_training.jl, loss/training_method_losses.jl, reduced_system/reduced_system.jl, the Hamiltonian, Lagrangian, symplectic-autoencoder, PSD and SAE-versus-PSD architecture tests, and optimizers/adam_with_learning_rate_decay.jl and optimizers/optimizer.jl — pass on Julia 1.13.2, exit 0: 1590 pass, 0 fail, 0 error, 1 broken (test/quality/jet.jl:111, #325). The full suite, the doctests and the docs build are CI's.

Findings 8 and 9 (the code half) stay unfixed, by the author's decision.

@michakraus
michakraus merged commit 37cae78 into main Oct 16, 2026
15 of 16 checks passed
@michakraus
michakraus deleted the simplify/part-m2 branch October 16, 2026 14:55
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Development

Successfully merging this pull request may close these issues.

Possible type instability in definition of DataLoader

1 participant