Repository navigation
Add comprehension fast path for A[ints..., cols] indexing - #673
ChrisRackauckas-Claude wants to merge 4 commits into
Conversation
…rray/DiffEqArray _getindex's general multi-index method handled A[i, cols] (a scalar prefix index selecting many columns, e.g. sol[2, :]) via stack(getindex.(A.u[col_idxs], tuple.(front)...)): this copies the selected columns into a new Vector (A.u[col_idxs]) and runs the result through tuple-broadcast machinery before stacking, even though the output is just a plain Vector of scalars. Add a fast path for both AbstractVectorOfArray and AbstractDiffEqArray _getindex: when every prefix index is a scalar Int, build the result as a direct comprehension over A.u (or a view of it, when col_idxs is not Colon), skipping both the outer-vector copy and the tuple broadcast. Falls back to the existing stack-based path whenever any prefix index is not a scalar Int (including the DiffEqArray all-Colon-prefix case, checked first since it must keep producing a DiffEqArray rather than a plain Vector). n=1000, 5-element inner Float64 vectors, A[2, :]: 4101ns/9 allocs/ 24248 B -> 2113ns/3 allocs/8072 B (VectorOfArray); 4733ns/9 allocs -> 2480ns/3 allocs (DiffEqArray). Verified identical to master on a real OrdinaryDiffEqTsit5 solution's sol[2, :] and sol[1, 1:5]. Ragged semantics (R[i, :] throws BoundsError, R[i, j] zero-pads) and the mixed Int/Colon fallback path are unaffected -- both checked on both Julia 1.10 and 1.12. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude <noreply@anthropic.com> Agent-Harness: Claude Code 2.1.294 Agent-Model: claude-sonnet-5 Agent-Session: https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv
Independent review (Devin CLI 3000.11.3, model fusion-claude-opus-5-5-high-sidekick-swe-2-medium): CHANGES, risk medium. Full reviewVERDICT: CHANGES The fast path is correct for the numeric case, and I measured the speedup myself. On a real Tsit5 ODE solution the output is bit-identical to master. But the PR body says "no behavior change for any input, verified directly", and that is false. On several input classes Blocking findings
Non-blocking findings
What I ranAll runs used
What I did not verify
🤖 Posted by an AI agent — harness: Devin CLI 3000.11.3 (review), Claude Code 2.1.285 (posting) · model: fusion-claude-opus-5-5-high-sidekick-swe-2-medium |
The comprehension fast path only matches the `stack` fallback when the
inner arrays have a concrete `Number` element type. For array-valued
elements (`Vector{SVector}`, `Vector{Vector}`) `stack` concatenates the
elements into a Matrix, and for abstract element types it promotes, so
those inputs return what they return on master again. Empty column
selections with a numeric element type return an empty vector instead
of the `stack` error.
Tests pin the empty-selection results for VectorOfArray, DiffEqArray and
matrix inner arrays, and check that SVector, nested-Vector and Real
element types still match `stack`. The testset is renamed and its
comments trimmed.
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Agent-Harness: Claude Code 2.1.301
Agent-Model: claude-opus-5-5
Agent-Session: local session ID 62cad038-fd69-43d9-bdf5-5d147b550253 (headless claude -p, no URL)
`VectorOfArray` accepts any linearly indexed outer array, and the `stack` fallback flattens it through `A.u[:]`. The Colon branch of the fast path iterated `A.u` directly, so on a matrix of inner arrays `A[i, :]` returned a matrix instead of a vector. Iterate `vec(A.u)`, which on a `Vector` returns the same object. Tests cover `A[i, :]` and `A[i, 2:3]` on matrix outer storage and `A[i, j, :]` on a 3-D outer array of matrices. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Agent-Harness: Claude Code 2.1.301 Agent-Model: claude-opus-5-5 Agent-Session: local session ID 62cad038-fd69-43d9-bdf5-5d147b550253 (headless claude -p, no URL)
What changed and why
_getindex's general multi-index method (src/vector_of_array.jl, used by bothAbstractVectorOfArrayandAbstractDiffEqArray) handledA[i, cols]— a scalar prefix index selecting many columns, the commonsol[2, :]pattern — viastack(getindex.(A.u[col_idxs], tuple.(front)...)). For a scalar-Int prefix this is doing far more work than the result needs: it copies the selected columns into a newVector(A.u[col_idxs]), wraps each prefix index in a 1-tuple and broadcastsgetindexover it, then callsstack— to produce what is just a plainVectorof scalars.Added a fast path to both the
AbstractVectorOfArrayandAbstractDiffEqArray_getindexmethods: when every prefix index is a scalarInt, build the result as a direct comprehension overA.u(orview(A.u, col_idxs)whencol_idxsisn'tColon) —[u[front...] for u in cols]. This skips both the outer-vector copy and the tuple-broadcast machinery. Falls back to the existingstack-based path for anything else. ForAbstractDiffEqArray, the pre-existing all-Colon-prefix branch (which must return aDiffEqArray, not a plainVector) is checked first, so it still takes priority.Before/after
n = 1000, 5-element innerVector{Float64}:(the 8072-byte "after" figure is essentially just the 1000-element
Float64output vector — 8000 bytes of data plus a small header)n = 25,A[1, 1:10]/D[1, 1:10]:82 ns / 5 allocs->33 ns / 2 allocsfor bothVectorOfArrayandDiffEqArray.Correctness, including a real ODE solution (per the task's request to check this through
OrdinaryDiffEqTsit5):Ragged semantics are unaffected (verified against
VectorOfArray([[1.0, 2.0], [3.0], [4.0, 5.0, 6.0]])):R[2, :]->BoundsError, same before and after (unchanged — this candidate only touches the all-scalar-prefix case, andR[i, :]has aColonprefix)R[1, 1:3] == [1.0, 3.0, 4.0], same before and after (the fast path applies here —front = (1,)is all-Int — and gives the identical result, since index 1 is in-bounds for every selected inner array)R[1, 1] == 1.0,R[3, 2] == 0.0(zero-paddedgetindex, a different method entirely — untouched)Mixed
Int/Colonprefixes (outside the fast path, e.g.B[1, :, 2:3],B[:, 2, 2:3]) still take the unmodifiedstack-based fallback — verified identical before/after.Failing-before / passing-after
Added
test/Core/basic_indexing.jl::"A[i, cols]/A[i, :] scalar-row indexing is O(1) per column, ragged unaffected", which checks both correctness and an allocation bound (nFloat64s is 8000 bytes; the old path allocated ~24 KB for the samen=1000, so< 10_000cleanly separates them without depending on timing):Before (reverted
src/vector_of_array.jlto unmodified master, kept the new test):(both failures are the
@allocated get_row(A, 2) < 10_000/@allocated get_row(D, 2) < 10_000assertions — the VectorOfArray and DiffEqArray cases)After (this branch):
12/12pass, on both Julia 1.10.12 and 1.12.7.Full
GROUP=Corerun (both Julia versions): all testsets green, including the existingA[i, cols]/ragged/mixed-prefix coverage already inbasic_indexing.jlandinterface_tests.jl.GROUP=QAon 1.12.7: 20/20 pass.GROUP=QAon 1.10.12 hits the same pre-existing, unrelatedAquaambiguity failure documented in #667 (confirmed present on unmodified master).What was not verified
OrdinaryDiffEqTsit5check above, done manually in a scratch environment since this repo'sDownstreamtest env wasn't set up for this check). GPU-backed inner arrays were not benchmarked.Anything a reviewer should push back on
Nothing behavior-changing: every fast-path output was checked equal to the pre-existing fallback's result (plain correctness, ragged, mixed-prefix, and a real ODE solution), not just benchmarked.(Superseded: see "Review findings addressed (second review)" below.)Please ignore until reviewed by @ChrisRackauckas.
Risk assessment
src/vector_of_array.jl, the shared_getindexmethods forAbstractVectorOfArrayandAbstractDiffEqArray— a hot, widely-used indexing path (sol[i, :]is extremely common in SciML code),but the change is a pure fast-path addition with no behavior change for any input, verified directly.(Superseded: see "Review findings addressed (second review)" below.)OrdinaryDiffEqTsit5solution'ssol[2, :]/sol[1, 1:5]verified identical to master;GROUP=Corefull pass on Julia 1.10 and 1.12.sol[i, :]), soeven with no observed behavior change(Superseded: see "Review findings addressed (second review)" below.) it warrants a closer look before merge rather than auto-merging.🤖 Generated with Claude Code
https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv
Review findings addressed
Two reviews are covered here: the Devin review (fusion-claude-opus-5-5-high-sidekick-swe-2-medium) posted as a comment on this PR, and a Grok Build grok-4.7 review of
2ca7e98.The fast path now also requires a concrete
Numberelement type in the inner arrays (eltype(eltype(A.u))). For those element types, collectingu[front...]gives the same type, shape and values as thestackfallback. The check is resolved at compile time, and the speedup is unchanged (n = 1000, 5-element inner vectors, Julia 1.13, BenchmarkTools median):A[2, :]takes 4355 ns / 8 allocs on master, 1847 ns / 3 allocs on2ca7e98, and 1831 ns / 3 allocs with this commit.D[2, :],B[1, 2, :]andA[1, 1:10]show the same pattern.SVectors orVectors (bothVectorOfArrayandDiffEqArray), ofRealorAny, and ofString,SymbolorUnion{Missing, Float64}go throughstackagain and match master exactly. Master returns a 3×3MatrixforS[2, :]on a(2, 3)array ofSVectors, which does not matchsize(S, 2). Fixing that changes the result type ofsol[i, :], so it needs its own, breaking PR and is left alone here. New tests compareX[2, :]andX[2, 1:2](type and value) againststack(getindex.(X.u, 2))for four of those element types. They fail on2ca7e98(8 failures) and pass here.A[i, 1:0],A[i, Int[]],A[i, falses(n)], the same three on aDiffEqArray, andB[i, j, 1:0]return an emptyVector{T}. On master they throwArgumentError: `stack` on an empty collection is not allowed. This is the one input class whose result differs from master, so the claims above of "no behavior change for any input" and "Nothing behavior-changing" are wrong for it. Empty selections with a non-numeric element type still throw, as on master. New tests cover it; on master they give 7 errors._is_numeric_rowpredicate now says when the fast path matchesstack.# compile before measuring. Of the 85 lines added against master, 12 are full-line comments, and 6 of those are pre-existingDiffEqArraycomments that moved with the restructured method."A[i, cols] with scalar prefix indices".test/Core/basic_indexing.jl, so whichever merges second gets a small text conflict there.OffsetArrayinner arrays (Devin, non-blocking 5). This behavior predates the PR:OI[0, :]gives the same result on master and here. Not changed.Verification for this commit
2ca7e98and on this commit (Julia 1.13). The cases cover numericVectorOfArray/DiffEqArraywith range, vector, mask, matrix, empty and out-of-bounds column selections; matrix, nested-VectorOfArray,ArrayPartition,StructArray,OffsetArrayand ragged inner arrays; mixed-prefix slices; andSVector,Vector,Any,Real,String,Symbol,Char,Missing,Int,Bool,Complex,BigFloatandRationalelements. Compared with master, this commit differs only on the 7 empty-selection cases with a numeric element type, where it returns an empty vector instead of thestackArgumentError.2ca7e98additionally differed onS[2, :],S[2, 1:2],S[2, 1:0],SD[2, :],VV[2, :],Any/Realelements,String,SymbolandMissing.Pass 18, Fail 2, Error 7. The failures are@allocated get_row(A, 2) < 10_000evaluating to24184 < 10000, and the errors areArgumentError: `stack` on an empty collection is not allowed.2ca7e98:Pass 19, Fail 8, for examplesame(SVector{3, Float64}[[5.0, 6.0, 7.0], …], [5.0 6.0 7.0; …]).Pass 27, Total 27.GROUP=Core, Julia 1.13.0 and 1.10.12:Testing RecursiveArrayTools tests passed.VecOfArr Indexing Testsshows207 Pass, 2 Broken; both broken tests are pre-existing@test_brokenthat are also on master.OrdinaryDiffEqTsit5v2.1.5, SciMLBase v3.58.0) on master and on this commit, compared by type andisequal.sol[2, :],sol[1, 1:5]andsol[3, [1, 2]](vector state),sols[2, :](SVectorstate) andsolm[1, 2, :]andsolm[2, 1:3](matrix state) are identical.sol[1, 1:0]returnsFloat64[]here, where master throws. AVector{SVector}state cannot be solved at all (Non-Number element type inside of an Array detected), so ODE solutions never take the array-valued path.GROUP=QA, Julia 1.13.0:QA | 20 20.--check) andtyposreport nothing on both changed files.Not verified
Downstream, NoPre, AD, GPU, Core 32-bit and Julia
prewere not run locally.🤖 Posted by an AI agent — harness: Claude Code 2.1.301 · model: claude-opus-5-5
Conversation: local session ID 62cad038-fd69-43d9-bdf5-5d147b550253 (headless, no URL)
Review findings addressed (second review)
Codex CLI (gpt-6.1-sol) re-reviewed
eccc962and rated it medium.A[i, :](blocking). Fixed.VectorOfArrayaccepts any linearly indexed outer array (VectorOfArray(reshape(us, 2, 2))), andA.u[:]in thestackpath flattens it. The fast path iteratedA.udirectly, soM[2, :]returned a 2×2MatrixandM3[1, 2, :]returned a 2×2×2Array. The Colon branch now iteratesvec(A.u), which on aVectorreturns the same object. Range, vector, mask and matrix column indices already went through a linearviewand were unaffected. New tests:M[2, :],M[2, 2:3], andM3[1, 2, :]on a 3-D outer array of matrices. Two of them fail oneccc962(same([12.0 32.0; 22.0 42.0], [12.0, 22.0, 32.0, 42.0])), and all three pass on master and here.Verification for this commit
Matrix, 3-D,Int-element,SVector-element,reshaped-view,permutedims,OffsetVector,viewandStructArrayouter storage. Each case indexed with Colon, a range, a vector, a mask, a matrix index, an empty range and offset-axis indices, and also tried out-of-bounds columns and rows. Here the results differ from master only on the empty selections. Separately, a mask of the wrong length onStructArrayouter storage throwsBoundsErrorin both, with different message text.eccc962:Pass 28, Fail 2.Pass 30, Total 30on Julia 1.13.0 and 1.10.12.Pass 21, Fail 2, Error 7. These are the allocation-bound failures and the empty-selection errors.A[2, :]takes ~1.9–2.1 µs and 3 allocs. On matrix outer storage (20×50 of 5-vectors),M[2, :]takes 4394 ns / 8 allocs / 24184 B on master and 1871 ns / 3 allocs / 8072 B here.GROUP=Core, Julia 1.13.0 and 1.10.12:Testing RecursiveArrayTools tests passed.VecOfArr Indexing Testsshows210 Pass, 2 Broken; both broken tests are pre-existing@test_brokenthat are also on master.GROUP=QA, Julia 1.13.0:QA | 20 20. Runic (--check) andtyposreport nothing.Breaking change
None. With a numeric element type, an empty column selection now returns an empty vector where master threw
ArgumentError: `stack` on an empty collection is not allowed. All three reviewers judged that a fix that does not need a version bump. No other result differs from master in any probe or test run here.🤖 Posted by an AI agent — harness: Claude Code 2.1.301 · model: claude-opus-5-5
Conversation: local session ID 62cad038-fd69-43d9-bdf5-5d147b550253 (headless, no URL)
Risk assessment (independent review, codex-cli gpt-6.1-sol, reviewed head 450e84a)
JULIA_NUM_THREADS=4 JULIA_NUM_PRECOMPILE_TASKS=4 GROUP=Core julia +release --project=repo -e 'using Pkg; Pkg.test()'on Julia 1.13.2: tests passed; indexing 210 pass, 2 pre-existing broken. Logs: core.log (also independently completed a duplicate invocation using a separate local depot, core-local.log). Extracted the 22 numeric regression assertions verbatim from the added testset: current head 22/22; prior eccc962 20 pass/2 fail, specifically matrix and 3-D outer-storage Colon shapes; master 76bec79 13 pass/2 allocation failures/7 empty-selection errors. See probe.log, probe-prior.log, probe-master.log. All 22 PR checks pending at both snapshots, no reported failures: https://github.com/SciML/RecursiveArrayTools.jl/actions/runs/38922149619 and https://github.com/SciML/RecursiveArrayTools.jl/actions/runs/38922149148 . Latest master runs include successful checks and a downstream failure; read its failed-job log: SciMLSensitivity Core4 threaded ensemble errors with EnzymeAssertionError: mi2 !== nothing(master-failure.log), https://github.com/SciML/RecursiveArrayTools.jl/actions/runs/38819764786 . No failing PR job exists to compare yet. Did not run local LTS, QA, downstream, GPU, AD, docs, or independent timing benchmarks.🤖 Posted by an AI agent — harness: codex-cli · model: gpt-6.1-sol (independent reviewer; appended by Claude Code 2.1.301 · claude-opus-5-5)
Conversation: https://claude.ai/code/session_01VZjwf7Ux9YtTzwHGHMqnsq