Repository navigation
GML part M2: DataLoader samples against trajectories, one count and one gather - #346
Conversation
…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 Report❌ Patch coverage is
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. 🚀 New features to boost your workflow:
|
michakraus
left a comment
There was a problem hiding this comment.
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
endThe 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)
endThen :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.jlandtest/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)GeometricSolutioninfer concrete types, with the kind fixed by a literal keyword. - The comment at
data_loader.jl:350-351holds: KernelAbstractions definescopyto!forCPUalone (src/cpu.jl:19), and JLArrays defines onlyBase.copyto!. - The deleted
rrule:test/data_loader/accuracy.jl:67-75checks Zygote's adjoint, zero on the cropped columns and with the primal's element type. - JuliaFormatter (
sciml) passes on the 7 changedsrc/testfiles, andtest-layout.jl --checkis 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
GMLDatasetscaller thatdata_loader.jl:51-53and 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]>
|
Follow-up to the review above. The fixes are in 6342d90. Every finding by its number:
Checks on 6342d90: JuliaFormatter and CI before this push: the one red job, Downgrade - ubuntu-latest, fails in the resolver ( |
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]>
|
Update on the two findings that were open in the follow-up above, both decided by the author:
Checks on 81bdfa8: the pre-commit hook (layout, format, lint, NFC, load) passed. 28 test files — every file in Findings 8 and 9 (the code half) stay unfixed, by the author's decision. |
The
DataLoaderhalf 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 itsformer size.
Fixes #211
Relates to #212: the batch indices reach the backend in one copy, and the
@viewsit asked forare declined (see One gather below).
Refs #325 (the batch half of the JET dispatch reports; the
PoissonTensorhalf is part M6's andthis change uses no closing keyword for it)
What this changes
DataLoader(input, output)holds independent samples, whatever its kind. At thebase it was typed
:TimeSeriesand drew the samples 1 to 9 of ten pairs, and where the input andthe output differed in array type one minibatch held the whole data set. It is now
:RegularDataand 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:TimeSeriesgives the input at offsetseq_length, and no output with:RegularDatagives theinput window alone.
1:nwithn = min(input_time_steps − seq_length + 1, target_steps − offset − prediction_window + 1), and the batch count iscld(n · n_params, batch_size).n < 1raises anArgumentErrornaming the time steps,seq_lengthandprediction_window, where the base asserted. The draw order is unchanged, so aseeded trajectory run draws the minibatches it drew before, except the ones the last sample joins.
Batch(batch_size)needs one time-axis length. ABatch{:FeedForward}pairs the input and theoutput 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
ArgumentErrorthat namesBatch(batch_size, dl), where it would train each sample on its first patch.Batch(batch_size, dl)is aBatch{:Transformer}, sotrain!on aNeuralNetworkIntegratorwith it needs anexplicit
loss.convert_input_and_batch_indices_to_arraymethods and four kernels becomeone
@kernel,out[i, j, k] = data[i, idx[1, k] + offset + j − 1, idx[2, k]], and one helper thatapplies it to an array or to a
q/ppair; the types ofoutanddataare independent, becausea view output is a
ReshapedArrayof aSubArray. The index matrix is one host-built2 × nMatrix{Int}copied to the backend once, and there are no@views— a list of(t, p)windows isnot a strided view, which is what Don't allocate arrays in
convert_input_and_batch_indices_to_arraybut use@views. #212 had asked for.autoencoderdefault. It isfalsein every constructor that takes a data set, with thematrix and vector methods'
truegone, andsuppress_infois a keyword of every one of thoseconstructors. The copying constructor
DataLoader(dl, backend, T)keepsautoencoder = nothing,meaning inherit. The manual, the transformer tutorial and the reproduction script that read a
matrix as a set of samples now pass
autoencoder = trueexplicitly.DataLoader#211.output_dimandoutput_time_stepscarry a fifth type parameterS <: Union{Int, Nothing};Kstays the fourth, soDataLoader{T, AT, OT, :TimeSeries}patterns still match. Afour-parameter constructor
DataLoader{T, AT, OT, K}(…)fillsSin from the fields, and a kernellaunch no longer reads a value of no single type.
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 andaccuracyis reachable through a public constructor.detect_ambiguitiesfinds none, anInt/Intpair builds
:RegularData, and aFloat32input with aFloat64output stays aMethodError.tensor_assign.jl.assign_output_estimate(x, pw) = x[:, (end − pw + 1):end, :]with no rrule:Zygote's
getindexadjoint is zero on the cropped columns, runs on a device with scalar indexingoff, and gives a
Float32cotangent for aFloat64seed onFloat32data, whereaugment_zerosreturned the seed's type. The kernel,
augment_zerosand the file-top string Julia discards go.DataLoader(::EnsembleSolution{T, T1, Vector{ST}})method for adataserof(t, q, q̇)is reachable — a rigid-bodyodeensembleproduces one, and the tutorialand three scripts call it — so it is not deleted;
test/data_loader/data_loader.jltests it.JET.report_optline oftest/quality/jet.jlfor a launcher that readsa
Union{Int, Nothing}field is a@testrather 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:
batch.jldata_loader.jltensor_assign.jl270 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.jland thedoctestsgroup — 1970 assertions, 0 failures.The four whole-package runs, all green at
a5e959e0, whose diff to the critic's head07ab13e9iscomment lines only with an identical code and docstring AST:
Testing GeometricMachineLearning tests passed, the only non-pass being JET's one known brokenPoissonTensortestjulia --project=docs docs/make.jl)training.jlwithGML_SMOKE=1The head carries, after those runs, the pre-PR verifier's fixes (an inference fix on the tensor
constructors, five comment corrections, the
KNOWN_ISSUES.mdK18 entry) and the correction of theCHANGELOG 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.
origin/mainaccuracy.jlheadersrc,test, accuracy headercalled it.remains in the baseline accuracy comment; replacement count 0test/data_loader/accuracy.jl; 1/1data_loader/accuracy.jlFloat32/Float64inputs; same-type,IntandBooltargets; mismatched float typesNeuralNetworkforwarding methodPer-file test counts of that round:
data_loader/accuracy.jlquality/reachability.jlquality/exports.jlquality/aqua.jlquality/jet.jlquality/docs_make_checks_references.jlEvery defect of the earlier rounds, as the same round reported it:
JLArraybatching fails before gatherallowscalar(false)pprecision insufficiently tested(q, p)gathered columns uncheckedFloat32target promotion survivescalled it.historyBlocking defects: None. Other defects: None. Fix first: none.
Pre-PR verification
Verdict: clear to open — 0 unresolved.
Verified
simplify/part-m2at07ab13e9, including its pending working-tree changes, against mergebase
d751ecd9.Base.@constprop :aggressivealone,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_infoclaim.
"without" replacement; two docstrings misdescribe index types or unconditional allocation. These
remain recorded and do not block this change.
batching tests 300/300, and index-allocation regressions 2/2 on both Julia 1.13.2 and the 1.12.7
floor. Live
code_typedchecks return concrete constructor and CPU minibatch types. Aqua piracypasses. 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.
JET/import sweep, solution/ensemble-specific inference probe, GPU measurement or timing benchmark
was performed.
Known issues added
KNOWN_ISSUES.mdgains one entry, K18 — the tensor-sampling formula in theDataLoadermanualis 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 takesn = input_time_steps − seq_length − prediction_window + 1and countscld(n · n_params, batch_size)batches (three start positions against the code's four for the page's five time stepsand
Batch(4, 2, 0)). It is a docs page the part's own text pass does not own, so it is recordedrather 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 (
accuracyis reachable through thepublic two-element-type constructor), K8 (a loader of one sample batches) and K12 (
Batchnow 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:
docs/src/data_loader/data_loader.md:143+ 1. The base has the same formula; recorded as K18.docs/src/data_loader/data_loader.md:98src/data_loader/batch.jl:12Vector{Vector{Tuple{Int64,Int64}}}. The sentence is outside the changed lines.src/data_loader/data_loader.jl:363