[CarloLucibello]

Closes #860.

The pullbacks of reduce(vcat, As), reduce(hcat, As) (reduce_vcat_pullback, reduce_hcat_pullback_1, reduce_hcat_pullback_2) and hvcat (hvcat_pullback) index the cotangent once per input. Indexing an AbstractThunk calls unthunk each time, so a thunked cotangent (e.g. from sum's pullback in Zygote, FluxML/Zygote.jl#1686) is recomputed length(As) times. This calls unthunk once at the top of each, as hcat/vcat/cat/stack already do.

I used dY = unthunk(ȳ) (renaming the argument) rather than dY = unthunk(dY), matching the neighbouring pullbacks and avoiding rebinding a variable captured by the do blocks. hvcat_pullback has the same pattern and is included.

Measured on Julia 1.13.2 (min of 5 runs), thunked cotangent @thunk(ones(size(Y))):

pullback	before	after	plain array
reduce(vcat), 1600 × Vector(10)	3.9 ms, 196 MiB	0.026 ms, 0.35 MiB	0.024 ms
same, InplaceableThunk	5.8 ms, 196 MiB	0.030 ms, 0.35 MiB	0.023 ms
reduce(vcat), 800 × Matrix(10×2)	1.5 ms, 98 MiB	0.026 ms, 0.31 MiB	0.017 ms
reduce(hcat), 1600 × Vector(10)	4.3 ms, 196 MiB	0.024 ms, 0.35 MiB	0.024 ms
reduce(hcat), 800 × Matrix(10×2)	1.6 ms, 98 MiB	0.025 ms, 0.31 MiB	0.020 ms
hvcat, 200 × Vector(10)	0.145 ms, 3.1 MiB	0.010 ms, 0.05 MiB	0.003 ms

Added a testset in test/rulesets/Base/array.jl checking that @thunk and InplaceableThunk cotangents give the same result as the plain array for all the variants above (vectors / matrices / mixed), and that @allocated back(thunk) < 10 * sizeof(Y) at 1600 inputs (the quadratic path allocates about length(As) * sizeof(Y)). Without the fix, the 10 allocation checks fail and the 14 equality checks pass; with it all 24 pass, also under --check-bounds=yes --code-coverage=user.

Tests: test/runtests.jl rulesets/Base on Julia 1.13.2 passes (16978 pass, 52 broken as before). Not run locally on 1.10 (left to CI).

Bumps the patch version to 1.74.10, as recent PRs do.

Noticed but left alone:

• vect pullbacks index ȳ[n] per element and are equally quadratic in the number of arguments (N=256: 0.8 ms vs 3 µs), but N is a compile-time argument count, so I did not touch them.
• Other pullbacks I looked at in src/rulesets/Base/ either already unthunk or access the cotangent once.

🧔‍♂️🤖 Assisted by AI

---

Follow-up edits (after review)

Addressing the review in 72b132c, 5f925a1 and 975805c:

• vect: the projected (Vararg{Union{Number,AbstractArray{<:Number}}}) and Vararg{Any} pullbacks now call unthunk once instead of indexing ȳ[n] per argument. The uniform Vararg{T} method converts with NTuple{N}(ȳ), which iterates and so unthunks only once, and is unchanged. Numeric arguments of a single type, e.g. vect(1.0, 2.0), dispatch to the projected method. Closes Rule for vect unthunks many times #676
• hvcat: hvcat_pullback returns one InplaceableThunk per argument, like hcat/vcat/cat, so slices for arguments that don't need a tangent are never computed. Its index computation uses d <= 2 instead of d in (1, 2), which Julia 1.10 doesn't constant-fold, so the returned thunks infer concretely there too (975805c).
• reduce(vcat/hcat) projection: the pullbacks now project each slice with ProjectTo.(As), as the hcat/vcat/cat pullbacks do. Before, a Vector input among matrices got an n×1 Matrix cotangent, and the element type wasn't projected (e.g. complex cotangents for Float32 inputs). The projectors are built by broadcasting: map(ProjectTo, As) doesn't infer concretely, and map over the existing adjs test input (a reshaped Adjoint) calls adjoint on the stored results. The two test_rrule(reduce, *, mix) calls on Vector{AbstractVecOrMat} inputs now use check_inferred=false. Their correctly typed cotangents (one Matrix, one Vector) can't be inferred concretely. Homogeneous inputs still pass the inference checks.
• Tests:
  • The @allocated threshold testset (1600 inputs, 200-argument hvcat) and the plain/thunk/inplaceable/allocs helpers are replaced by a testset that counts how often a @thunk / InplaceableThunk cotangent is evaluated (exactly once) and compares unthunked results against the plain-array cotangent. It covers the 7 reduce/hvcat cases above plus the three vect methods. Against the source on main, 18 of its 40 checks fail (every count check except the uniform vect ones).
  • A new testset reduce cat: cotangents are projected checks that the cotangents of Float32 vector/matrix inputs keep the inputs' types for a ComplexF64 cotangent. All 3 of its checks fail before 5f925a1.
  • All of these, plus the existing vect, hvcat, reduce hcat and reduce vcat testsets, pass locally (Julia 1.13.2).

Updated benchmarks

Pullback back(dȲ) on main (a441bd5) and this PR, Julia 1.13.1 (aarch64, macOS), minimum time and memory from BenchmarkTools. @thunk is @thunk(ones(size(Y))) and plain is ones(size(Y)). InplaceableThunk(dx -> dx .+= 1, @thunk(ones(size(Y)))) cotangents give the same numbers as @thunk to within noise.

pullback	@thunk, main	@thunk, PR	plain, main	plain, PR
reduce(vcat), 1600 × Vector(10)	4.06 ms, 200.3 MiB	21.0 µs, 367.1 KiB	18.2 µs, 239.1 KiB	18.2 µs, 239.1 KiB
reduce(vcat), 800 × Matrix(10×2)	1.80 ms, 100.3 MiB	17.0 µs, 322.7 KiB	14.0 µs, 194.6 KiB	15.6 µs, 194.6 KiB
reduce(hcat), 1600 × Vector(10)	3.86 ms, 200.4 MiB	18.0 µs, 367.2 KiB	16.0 µs, 239.1 KiB	17.1 µs, 239.1 KiB
reduce(hcat), 800 × Matrix(10×2)	1.92 ms, 100.3 MiB	17.2 µs, 322.7 KiB	14.3 µs, 194.6 KiB	15.1 µs, 194.6 KiB
reduce(hcat), 10 × Vector(10)	1.1 µs, 11.1 KiB	160 ns, 2.5 KiB	129 ns, 1.6 KiB	129 ns, 1.6 KiB
hvcat, 200 × Vector(10)	70.7 µs, 3.2 MiB	12.5 µs, 77.2 KiB	10.3 µs, 33.0 KiB	11.9 µs, 61.2 KiB
vect, 256 × Float64	33.7 µs, 542.3 KiB	17.2 µs, 16.4 KiB	16.4 µs, 14.3 KiB	17.1 µs, 14.3 KiB
vect, Symbol + 255 × Float64	34.6 µs, 542.3 KiB	16.5 µs, 16.4 KiB	16.4 µs, 14.3 KiB	16.4 µs, 14.3 KiB

• The ProjectTo calls added to reduce(vcat/hcat) cost at most about 1.5 µs with 800–1600 inputs and nothing measurable with 10.
• On this PR, hvcat returns lazy InplaceableThunks, so its rows measure back alone. Unthunking all 200 outputs takes about another 9 µs and 34 KiB, the same per-argument cost as the existing hcat rule.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.64%. Comparing base (a441bd5) to head (975805c).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #861      +/-   ##
==========================================
+ Coverage   96.63%   96.64%   +0.01%     
==========================================
  Files          28       28              
  Lines        3598     3609      +11     
==========================================
+ Hits         3477     3488      +11     
  Misses        121      121              

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

[devmotion]

Thanks, the fix looks good. Unthunking once at the top matches hcat/vcat/cat/stack. It leaves plain-array and ZeroTangent cotangents alone, and it handles InplaceableThunk.

Two suggestions on the tests:

1. Rather than using @allocated thresholds as a stand-in, consider counting unthunks directly. That checks the property itself, gives the same result on every Julia version and coverage setting, and doesn't need the 1600-input or 200-argument hvcat cases:

   @testset "$name" for (name, (Y, back)) in cases
       dY = rand(size(Y)...)
       expected = back(dY)
       n = Ref(0)
       counted() = (n[] += 1; copy(dY))
       for dȲ in (@thunk(counted()), InplaceableThunk(dx -> dx .+= dY, @thunk(counted())))
           n[] = 0
           @test back(dȲ) == expected
           @test n[] == 1
       end
   end

   This would replace the "cost does not grow with the number of inputs" testset and the allocs/thunk/inplaceable helpers.
2. plain(dY) = dY can go; expected = back(dY) does the same.

Possible follow-up: hvcat_pullback could return one InplaceableThunk per argument, as hcat/vcat/cat do, so slices for non-differentiable arguments are never computed.

[devmotion]

I pushed two commits that address my review and the vect case as well.

[CarloLucibello]

thanks!