From d32589bff1f1af33bb05e2399556259be3420c1e Mon Sep 17 00:00:00 2001 From: CarloLucibello Date: Sat, 31 Oct 2026 10:27:11 +0100 Subject: [PATCH 1/4] Unthunk the cotangent in `reduce(vcat/hcat)` and `hvcat` pullbacks These pullbacks index the cotangent once per input array. Indexing an `AbstractThunk` forwards to `unthunk`, which recomputes the whole array on every access, so the cost was O(length(As) * length(dY)). Unthunk once at the top of each pullback, as the other `cat` pullbacks already do. Also bumps the patch version. Co-Authored-By: Claude Sonnet 5.5 --- Project.toml | 2 +- src/rulesets/Base/array.jl | 12 ++++++---- test/rulesets/Base/array.jl | 48 +++++++++++++++++++++++++++++++++++++ 3 files changed, 57 insertions(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index a18c5c4b0..952b846d4 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ChainRules" uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2" -version = "1.74.9" +version = "1.74.10" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/rulesets/Base/array.jl b/src/rulesets/Base/array.jl index 7e4098509..8d17df419 100644 --- a/src/rulesets/Base/array.jl +++ b/src/rulesets/Base/array.jl @@ -259,7 +259,8 @@ function rrule(::typeof(reduce), ::typeof(hcat), As::AbstractVector{<:AbstractVe Y = reduce(hcat, As) Base.require_one_based_indexing(Y) widths = map(A -> size(A,2), As) - function reduce_hcat_pullback_2(dY) + function reduce_hcat_pullback_2(ȳ) + dY = unthunk(ȳ) hi = Ref(0) dAs = map(widths) do w lo = hi[]+1 @@ -273,7 +274,8 @@ end function rrule(::typeof(reduce), ::typeof(hcat), As::AbstractVector{<:AbstractVector}) axe = axes(As,1) - function reduce_hcat_pullback_1(dY) + function reduce_hcat_pullback_1(ȳ) + dY = unthunk(ȳ) hi = Ref(0) dAs = map(_ -> dY[:, hi[]+=1], axe) return (NoTangent(), NoTangent(), dAs) @@ -328,7 +330,8 @@ function rrule(::typeof(reduce), ::typeof(vcat), As::AbstractVector{<:AbstractVe Base.require_one_based_indexing(Y) ndimsY = Val(ndims(Y)) heights = map(A -> size(A,1), As) - function reduce_vcat_pullback(dY) + function reduce_vcat_pullback(ȳ) + dY = unthunk(ȳ) hi = Ref(0) dAs = map(heights) do z lo = hi[]+1 @@ -398,7 +401,8 @@ function rrule(::typeof(hvcat), rows, values...) ndimsY = Val(ndims(Y)) sizes = map(_catsize, values) project_Vs = map(ProjectTo, values) - function hvcat_pullback(dY) + function hvcat_pullback(ȳ) + dY = unthunk(ȳ) prev = fill(0, 2) dXs = map(project_Vs, sizes) do project, sizeX ndimsX = length(sizeX) diff --git a/test/rulesets/Base/array.jl b/test/rulesets/Base/array.jl index 46a52a0a8..f5c657365 100644 --- a/test/rulesets/Base/array.jl +++ b/test/rulesets/Base/array.jl @@ -277,6 +277,54 @@ end test_rrule(hvcat, (1,2), rand(2)', (3.4, 5.6), 7.8; check_inferred=false) end +@testset "reduce cat & hvcat: thunked cotangents" begin + # Unthunking recomputes the array, so the pullbacks must do so only once, not per input. + plain(dY) = dY + thunk(dY) = @thunk(copy(dY)) + inplaceable(dY) = InplaceableThunk(dx -> dx .+= dY, @thunk(copy(dY))) + allocs(back, dY) = (back(dY); @allocated back(dY)) + + cases = [ + ("reduce vcat, vectors", rrule(reduce, vcat, [rand(3), rand(2), rand(4)])), + ("reduce vcat, matrices", rrule(reduce, vcat, [rand(2, 4), rand(1, 4), rand(3, 4)])), + ("reduce vcat, mixed", rrule(reduce, vcat, AbstractVecOrMat[rand(4, 1), rand(4)])), + ("reduce hcat, vectors", rrule(reduce, hcat, [rand(3) for _ in 1:4])), + ("reduce hcat, matrices", rrule(reduce, hcat, [rand(3, 2), rand(3, 1), rand(3, 3)])), + ("reduce hcat, mixed", rrule(reduce, hcat, AbstractVecOrMat[rand(4, 2), rand(4)])), + ("hvcat", rrule(hvcat, (2, 1), rand(), rand(1, 1), rand(2, 2))), + ] + @testset "$name" for (name, (Y, back)) in cases + dY = rand(size(Y)...) + expected = back(plain(dY)) + @test back(thunk(dY)) == expected + @test back(inplaceable(dY)) == expected + end + + @testset "cost does not grow with the number of inputs" begin + n = 1600 + for (name, red, As) in [ + ("vcat, vectors", vcat, [rand(10) for _ in 1:n]), + ("vcat, matrices", vcat, [rand(10, 2) for _ in 1:n]), + ("hcat, vectors", hcat, [rand(10) for _ in 1:n]), + ("hcat, matrices", hcat, [rand(10, 2) for _ in 1:n]), + ] + Y, back = rrule(reduce, red, As) + dY = rand(size(Y)...) + # Unthunking once per input would allocate `n` times `sizeof(Y)`. + @testset "$name, $(nameof(mk))" for mk in (thunk, inplaceable) + @test allocs(back, mk(dY)) < 10 * sizeof(Y) + end + end + + As = [rand(10) for _ in 1:200] + Y, back = rrule(hvcat, (200,), As...) + dY = rand(size(Y)...) + @testset "hvcat, $(nameof(mk))" for mk in (thunk, inplaceable) + @test allocs(back, mk(dY)) < 10 * sizeof(Y) + end + end +end + @testset "reverse" begin @testset "Tuple" begin test_frule(reverse, Tuple(rand(10))) From 72b132c3ab7fc8edcd8862eead70a8cbe1328805 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 3 Nov 2026 12:28:53 +0100 Subject: [PATCH 2/4] Unthunk once in `vect` pullbacks, return thunks from `hvcat` pullback, count unthunks in tests - `vect` pullbacks for numbers/arrays and for arbitrary arguments indexed the cotangent once per argument, unthunking each time. - `hvcat_pullback` returns an `InplaceableThunk` per argument, as the `hcat`/`vcat`/`cat` pullbacks do. - Tests count unthunks directly instead of checking allocations. Co-Authored-By: Claude Opus 5.5 --- src/rulesets/Base/array.jl | 13 ++++++++--- test/rulesets/Base/array.jl | 43 ++++++++++--------------------------- 2 files changed, 21 insertions(+), 35 deletions(-) diff --git a/src/rulesets/Base/array.jl b/src/rulesets/Base/array.jl index 8d17df419..2497e5cc6 100644 --- a/src/rulesets/Base/array.jl +++ b/src/rulesets/Base/array.jl @@ -50,7 +50,8 @@ function rrule( ) where {N} projects = map(ProjectTo, X) function vect_pullback(ȳ) - X̄ = ntuple(n -> projects[n](ȳ[n]), N) + dY = unthunk(ȳ) + X̄ = ntuple(n -> projects[n](dY[n]), N) return (NoTangent(), X̄...) end return Base.vect(X...), vect_pullback @@ -58,7 +59,10 @@ end # Data is unmodified, so no need to project. function rrule(::typeof(Base.vect), X::Vararg{Any,N}) where {N} - vect_pullback(ȳ) = (NoTangent(), ntuple(n -> ȳ[n], N)...) + function vect_pullback(ȳ) + dY = unthunk(ȳ) + return (NoTangent(), ntuple(n -> dY[n], N)...) + end return Base.vect(X...), vect_pullback end @@ -418,7 +422,10 @@ function rrule(::typeof(hvcat), rows, values...) prev[2] = 0 prev[1] += get(sizeX, 1, 1) end - project(dY[index...]) + InplaceableThunk( + dX -> dX .+= view(dY, index...), + @thunk project(@allowscalar dY[index...]) + ) end return (NoTangent(), NoTangent(), dXs...) end diff --git a/test/rulesets/Base/array.jl b/test/rulesets/Base/array.jl index f5c657365..a205a1fea 100644 --- a/test/rulesets/Base/array.jl +++ b/test/rulesets/Base/array.jl @@ -277,13 +277,8 @@ end test_rrule(hvcat, (1,2), rand(2)', (3.4, 5.6), 7.8; check_inferred=false) end -@testset "reduce cat & hvcat: thunked cotangents" begin +@testset "thunked cotangents are unthunked once" begin # Unthunking recomputes the array, so the pullbacks must do so only once, not per input. - plain(dY) = dY - thunk(dY) = @thunk(copy(dY)) - inplaceable(dY) = InplaceableThunk(dx -> dx .+= dY, @thunk(copy(dY))) - allocs(back, dY) = (back(dY); @allocated back(dY)) - cases = [ ("reduce vcat, vectors", rrule(reduce, vcat, [rand(3), rand(2), rand(4)])), ("reduce vcat, matrices", rrule(reduce, vcat, [rand(2, 4), rand(1, 4), rand(3, 4)])), @@ -292,35 +287,19 @@ end ("reduce hcat, matrices", rrule(reduce, hcat, [rand(3, 2), rand(3, 1), rand(3, 3)])), ("reduce hcat, mixed", rrule(reduce, hcat, AbstractVecOrMat[rand(4, 2), rand(4)])), ("hvcat", rrule(hvcat, (2, 1), rand(), rand(1, 1), rand(2, 2))), + ("vect, uniform", rrule(Base.vect, (1.0,), (2.0,))), + ("vect, projected", rrule(Base.vect, 1.0, 2f0, 3)), + ("vect, any", rrule(Base.vect, :a, 1.0)), ] @testset "$name" for (name, (Y, back)) in cases dY = rand(size(Y)...) - expected = back(plain(dY)) - @test back(thunk(dY)) == expected - @test back(inplaceable(dY)) == expected - end - - @testset "cost does not grow with the number of inputs" begin - n = 1600 - for (name, red, As) in [ - ("vcat, vectors", vcat, [rand(10) for _ in 1:n]), - ("vcat, matrices", vcat, [rand(10, 2) for _ in 1:n]), - ("hcat, vectors", hcat, [rand(10) for _ in 1:n]), - ("hcat, matrices", hcat, [rand(10, 2) for _ in 1:n]), - ] - Y, back = rrule(reduce, red, As) - dY = rand(size(Y)...) - # Unthunking once per input would allocate `n` times `sizeof(Y)`. - @testset "$name, $(nameof(mk))" for mk in (thunk, inplaceable) - @test allocs(back, mk(dY)) < 10 * sizeof(Y) - end - end - - As = [rand(10) for _ in 1:200] - Y, back = rrule(hvcat, (200,), As...) - dY = rand(size(Y)...) - @testset "hvcat, $(nameof(mk))" for mk in (thunk, inplaceable) - @test allocs(back, mk(dY)) < 10 * sizeof(Y) + expected = map(unthunk, back(dY)) + n = Ref(0) + counted() = (n[] += 1; copy(dY)) + for dȲ in (@thunk(counted()), InplaceableThunk(dx -> dx .+= dY, @thunk(counted()))) + n[] = 0 + @test map(unthunk, back(dȲ)) == expected + @test n[] == 1 end end end From 5f925a14650a18129d5a8ee8fce6d96b8b134f41 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 3 Nov 2026 12:56:21 +0100 Subject: [PATCH 3/4] Project cotangents in `reduce(vcat/hcat)` pullbacks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A `Vector` input of `reduce(vcat/hcat)` with matrices received an `n×1` `Matrix` cotangent, and the cotangent's element type was not projected onto the inputs' (e.g. complex cotangents for real inputs). Project each slice with `ProjectTo`, as the `hcat`/`vcat`/`cat` pullbacks do. Co-Authored-By: Claude Opus 5.5 --- src/rulesets/Base/array.jl | 14 ++++++++------ test/rulesets/Base/array.jl | 15 +++++++++++++-- 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/src/rulesets/Base/array.jl b/src/rulesets/Base/array.jl index 2497e5cc6..2c5b31cd7 100644 --- a/src/rulesets/Base/array.jl +++ b/src/rulesets/Base/array.jl @@ -263,13 +263,14 @@ function rrule(::typeof(reduce), ::typeof(hcat), As::AbstractVector{<:AbstractVe Y = reduce(hcat, As) Base.require_one_based_indexing(Y) widths = map(A -> size(A,2), As) + project_As = ProjectTo.(As) function reduce_hcat_pullback_2(ȳ) dY = unthunk(ȳ) hi = Ref(0) - dAs = map(widths) do w + dAs = map(project_As, widths) do project, w lo = hi[]+1 hi[] += w - dY[:, lo:hi[]] + project(dY[:, lo:hi[]]) end return (NoTangent(), NoTangent(), dAs) end @@ -277,11 +278,11 @@ function rrule(::typeof(reduce), ::typeof(hcat), As::AbstractVector{<:AbstractVe end function rrule(::typeof(reduce), ::typeof(hcat), As::AbstractVector{<:AbstractVector}) - axe = axes(As,1) + project_As = ProjectTo.(As) function reduce_hcat_pullback_1(ȳ) dY = unthunk(ȳ) hi = Ref(0) - dAs = map(_ -> dY[:, hi[]+=1], axe) + dAs = map(project -> project(dY[:, hi[]+=1]), project_As) return (NoTangent(), NoTangent(), dAs) end return reduce(hcat, As), reduce_hcat_pullback_1 @@ -334,14 +335,15 @@ function rrule(::typeof(reduce), ::typeof(vcat), As::AbstractVector{<:AbstractVe Base.require_one_based_indexing(Y) ndimsY = Val(ndims(Y)) heights = map(A -> size(A,1), As) + project_As = ProjectTo.(As) function reduce_vcat_pullback(ȳ) dY = unthunk(ȳ) hi = Ref(0) - dAs = map(heights) do z + dAs = map(project_As, heights) do project, z lo = hi[]+1 hi[] += z ind = ntuple(d -> d==1 ? (lo:hi[]) : (:), ndimsY) - dY[ind...] + project(dY[ind...]) end return (NoTangent(), NoTangent(), dAs) end diff --git a/test/rulesets/Base/array.jl b/test/rulesets/Base/array.jl index a205a1fea..b94cf2600 100644 --- a/test/rulesets/Base/array.jl +++ b/test/rulesets/Base/array.jl @@ -199,7 +199,7 @@ end test_rrule(reduce, hcat, vecs) mix = AbstractVecOrMat[rand(4,2), rand(4)] # this is weird, but does hit the fast path - test_rrule(reduce, hcat, mix) + test_rrule(reduce, hcat, mix; check_inferred=false) adjs = vec([randn(2, 4), randn(1, 4), randn(3, 4)]') # not a Vector # test_rrule(reduce, hcat, adjs ⊢ map(m -> rand(size(m)), adjs)) @@ -240,7 +240,18 @@ end test_rrule(reduce, vcat, vecs) mix = AbstractVecOrMat[rand(4,1), rand(4)] - test_rrule(reduce, vcat, mix) + test_rrule(reduce, vcat, mix; check_inferred=false) +end + +@testset "reduce cat: cotangents are projected" begin + @testset "$red" for (red, As) in [ + (hcat, [rand(Float32, 3) for _ in 1:2]), + (hcat, AbstractVecOrMat[rand(Float32, 3, 2), rand(Float32, 3)]), + (vcat, AbstractVecOrMat[rand(Float32, 3, 1), rand(Float32, 3)]), + ] + Y, back = rrule(reduce, red, As) + @test map(typeof, back(rand(ComplexF64, size(Y)))[3]) == map(typeof, As) + end end @testset "cat" begin From 975805c279119801a865329c7a069101956b8b3e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 3 Nov 2026 13:49:28 +0100 Subject: [PATCH 4/4] Make `hvcat` pullback index inferable on Julia 1.10 `d in (1, 2)` is not constant-folded inside the `ntuple` closure on Julia 1.10, so the inferred index type is a `Union`. Since the pullback now returns thunks that capture the index, the `Union` leaked into its return type. `d <= 2` is folded, like the `d == 1`/`d == 2` checks in `vcat`/`hcat`. Co-Authored-By: Claude Opus 5.5 --- src/rulesets/Base/array.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/rulesets/Base/array.jl b/src/rulesets/Base/array.jl index 2c5b31cd7..9b12d1439 100644 --- a/src/rulesets/Base/array.jl +++ b/src/rulesets/Base/array.jl @@ -413,7 +413,7 @@ function rrule(::typeof(hvcat), rows, values...) dXs = map(project_Vs, sizes) do project, sizeX ndimsX = length(sizeX) index = ntuple(ndimsY) do d - if d in (1, 2) + if d <= 2 d > ndimsX ? (prev[d]+1) : (prev[d]+1:prev[d]+sizeX[d]) else d > ndimsX ? 1 : (:)