Port more array-related derivatives from Nabla (#45)

* reshape
* hcat
* vcat
* fill

Author: ararslan

src/ChainRules.jl

@@ -11,6 +11,7 @@ export AbstractRule, Rule, frule, rrule
 include("differentials.jl")
 include("rules.jl")
 include("rules/base.jl")
+include("rules/array.jl")
 include("rules/broadcast.jl")
 include("rules/linalg/dense.jl")
 include("rules/linalg/diagonal.jl")

src/rules/array.jl

@@ -0,0 +1,58 @@
+#####
+##### `reshape`
+#####
+
+function rrule(::typeof(reshape), A::AbstractArray, dims::Tuple{Vararg{Int}})
+    return reshape(A, dims), (Rule(Ȳ->reshape(Ȳ, dims)), DNERule())
+end
+
+function rrule(::typeof(reshape), A::AbstractArray, dims::Int...)
+    Y, (rule, _) = rrule(reshape, A, dims)
+    return Y, (rule, fill(DNERule(), length(dims))...)
+end
+
+#####
+##### `hcat` (🐈)
+#####
+
+function rrule(::typeof(hcat), A::AbstractArray, Bs::AbstractArray...)
+    Y = hcat(A, Bs...)
+    Xs = (A, Bs...)
+    rules = ntuple(length(Bs) + 1) do i
+        l = mapreduce(j->size(Xs[j], 2), Base.add_sum, 1:i-1; init=0)
+        u = l + size(Xs[i], 2)
+        dim = u > l + 1 ? (l+1:u) : u
+        # NOTE: The copy here is defensive, since `selectdim` returns a view which we can
+        # materialize with `copy`
+        Rule(Ȳ->copy(selectdim(Ȳ, 2, dim)))
+    end
+    return Y, rules
+end
+
+#####
+##### `vcat`
+#####
+
+function rrule(::typeof(vcat), A::AbstractArray, Bs::AbstractArray...)
+    Y = vcat(A, Bs...)
+    n = size(A, 1)
+    ∂A = Rule(Ȳ->copy(selectdim(Ȳ, 1, 1:n)))
+    ∂Bs = ntuple(length(Bs)) do i
+        l = n + mapreduce(j->size(Bs[j], 1), Base.add_sum, 1:i-1; init=0)
+        u = l + size(Bs[i], 1)
+        Rule(Ȳ->copy(selectdim(Ȳ, 1, l+1:u)))
+    end
+    return Y, (∂A, ∂Bs...)
+end
+
+#####
+##### `fill`
+#####
+
+function rrule(::typeof(fill), value::Any, dims::Tuple{Vararg{Int}})
+    return fill(value, dims), (Rule(sum), DNERule())
+end
+
+function rrule(::typeof(fill), value::Any, dims::Int...)
+    return fill(value, dims), (Rule(sum), ntuple(_->DNERule(), length(dims))...)
+end

test/rules/array.jl

@@ -0,0 +1,56 @@
+@testset "reshape" begin
+    rng = MersenneTwister(1)
+    A = randn(rng, 4, 5)
+    B, (dA, dd) = rrule(reshape, A, (5, 4))
+    @test B == reshape(A, (5, 4))
+    @test dd isa ChainRules.DNERule
+    Ȳ = randn(rng, 4, 5)
+    Ā = dA(Ȳ)
+    @test Ā == reshape(Ȳ, (5, 4))
+
+    B, (dA, dd1, dd2) = rrule(reshape, A, 5, 4)
+    @test B == reshape(A, 5, 4)
+    @test dd1 isa ChainRules.DNERule
+    @test dd2 isa ChainRules.DNERule
+    Ȳ = randn(rng, 4, 5)
+    Ā = dA(Ȳ)
+    @test Ā == reshape(Ȳ, 5, 4)
+end
+
+@testset "hcat" begin
+    rng = MersenneTwister(2)
+    A = randn(rng, 3, 2)
+    B = randn(rng, 3)
+    C = randn(rng, 3, 3)
+    H, (dA, dB, dC) = rrule(hcat, A, B, C)
+    @test H == hcat(A, B, C)
+    H̄ = randn(rng, 3, 6)
+    @test dA(H̄) ≈ view(H̄, :, 1:2)
+    @test dB(H̄) ≈ view(H̄, :, 3)
+    @test dC(H̄) ≈ view(H̄, :, 4:6)
+end
+
+@testset "vcat" begin
+    rng = MersenneTwister(3)
+    A = randn(rng, 2, 4)
+    B = randn(rng, 1, 4)
+    C = randn(rng, 3, 4)
+    V, (dA, dB, dC) = rrule(vcat, A, B, C)
+    @test V == vcat(A, B, C)
+    V̄ = randn(rng, 6, 4)
+    @test dA(V̄) ≈ view(V̄, 1:2, :)
+    @test dB(V̄) ≈ view(V̄, 3:3, :)
+    @test dC(V̄) ≈ view(V̄, 4:6, :)
+end
+
+@testset "fill" begin
+    y, (dv, dd) = rrule(fill, 44, 4)
+    @test y == [44, 44, 44, 44]
+    @test dd isa ChainRules.DNERule
+    @test dv(ones(Int, 4)) == 4
+
+    y, (dv, dd) = rrule(fill, 2.0, (3, 3, 3))
+    @test y == fill(2.0, (3, 3, 3))
+    @test dd isa ChainRules.DNERule
+    @test dv(ones(3, 3, 3)) ≈ 27.0
+end

test/runtests.jl

@@ -14,6 +14,7 @@ include("test_util.jl")
     include("rules.jl")
     @testset "rules" begin
         include(joinpath("rules", "base.jl"))
+        include(joinpath("rules", "array.jl"))
         @testset "linalg" begin
             include(joinpath("rules", "linalg", "dense.jl"))
             include(joinpath("rules", "linalg", "diagonal.jl"))