Unthunk the cotangent in reduce(vcat/hcat) and hvcat pullbacks

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"

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)

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)))