From d32589bff1f1af33bb05e2399556259be3420c1e Mon Sep 17 00:00:00 2001 From: CarloLucibello Date: Sat, 31 Oct 2026 10:27:11 +0100 Subject: [PATCH] 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)))