add adjoint for reduce hvcat (#236)

* add adjoint for reduce hvcat

* Apply suggestions from code review

Co-authored-by: Seth Axen <seth.axen@gmail.com>

* Change return type to DoesNotExist

Co-authored-by: Seth Axen <seth.axen@gmail.com>

* test with rrule_test

* use randn everywhere

* bump version

Co-authored-by: Seth Axen <seth.axen@gmail.com>

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "0.7.10"
+version = "0.7.11"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/rulesets/Base/array.jl

@@ -39,6 +39,19 @@ function rrule(::typeof(hcat), A::AbstractArray, Bs::AbstractArray...)
     return hcat(A, Bs...), hcat_pullback
 end
 
+function rrule(::typeof(reduce), ::typeof(hcat), As::AbstractVector{<:AbstractVecOrMat})
+    function reduce_hcat_pullback(ΔY)
+        sizes = size.(As, 2)
+        cumsizes = cumsum(sizes)
+        ∂As = map(cumsizes, sizes) do post, diff
+            pre = post - diff + 1
+            return ΔY[:, pre:post]
+        end
+        return (NO_FIELDS, DoesNotExist(), ∂As)
+    end
+    return reduce(hcat, As), reduce_hcat_pullback
+end
+
 #####
 ##### `vcat`
 #####
@@ -57,6 +70,19 @@ function rrule(::typeof(vcat), A::AbstractArray, Bs::AbstractArray...)
     return vcat(A, Bs...), vcat_pullback
 end
 
+function rrule(::typeof(reduce), ::typeof(vcat), As::AbstractVector{<:AbstractVecOrMat})
+    function reduce_vcat_pullback(ΔY)
+        sizes = size.(As, 1)
+        cumsizes = cumsum(sizes)
+        ∂As = map(cumsizes, sizes) do post, diff
+            pre = post - diff + 1
+            return ΔY[pre:post, :]
+        end
+        return (NO_FIELDS, DoesNotExist(), ∂As)
+    end
+    return reduce(vcat, As), reduce_vcat_pullback
+end
+
 #####
 ##### `fill`
 #####

test/rulesets/Base/array.jl

@@ -34,6 +34,18 @@ end
     @test dC ≈ view(H̄, :, 4:6)
 end
 
+@testset "reduce hcat" begin
+    A = randn(3, 2)
+    B = randn(3, 1)
+    C = randn(3, 3)
+    x = [A, B, C]
+    H, pullback = rrule(reduce, hcat, x)
+    @test H == reduce(hcat, x)
+    H̄ = randn(3, 6)
+    x̄ = randn.(size.(x))
+    rrule_test(reduce, H̄, (hcat, nothing), (x, x̄))
+end
+
 @testset "vcat" begin
     A = randn(2, 4)
     B = randn(1, 4)
@@ -48,6 +60,18 @@ end
     @test dC ≈ view(V̄, 4:6, :)
 end
 
+@testset "reduce vcat" begin
+    A = randn(2, 4)
+    B = randn(1, 4)
+    C = randn(3, 4)
+    x = [A, B, C]
+    V, pullback = rrule(reduce, vcat, x)
+    @test V == reduce(vcat, x)
+    V̄ = randn(6, 4)
+    x̄ = randn.(size.(x))
+    rrule_test(reduce, V̄, (vcat, nothing), (x, x̄))
+end
+
 @testset "fill" begin
     y, pullback = rrule(fill, 44, 4)
     @test y == [44, 44, 44, 44]