Unthunk the cotangent in reduce(vcat/hcat), hvcat and vect 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

@@ -50,15 +50,19 @@ 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
 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
 
@@ -259,23 +263,26 @@ 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)
+    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
     return Y, reduce_hcat_pullback_2
 end
 
 function rrule(::typeof(reduce), ::typeof(hcat), As::AbstractVector{<:AbstractVector})
-    axe = axes(As,1)
-    function reduce_hcat_pullback_1(dY)
+    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
@@ -328,13 +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)
-    function reduce_vcat_pullback(dY)
+    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
@@ -398,12 +407,13 @@ 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)
             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 : (:)
@@ -414,7 +424,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

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
@@ -277,6 +288,33 @@ end
     test_rrule(hvcat, (1,2), rand(2)', (3.4, 5.6), 7.8; check_inferred=false)
 end
 
+@testset "thunked cotangents are unthunked once" begin
+    # Unthunking recomputes the array, so the pullbacks must do so only once, not per input.
+    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))),
+        ("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 = 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
+
 @testset "reverse" begin
     @testset "Tuple" begin
         test_frule(reverse, Tuple(rand(10)))