Some Array rules (#491)

* Bump minor version

* Add undef Array constructor

* construct Array from existing AbstractArray

* vect implementation

* Bump precision

* Additional tests

Co-authored-by: Michael Abbott <[email protected]>

* Fix undef tests

* Constraint vect implementation

* Add float-only test

* Link to non_differentiable tests issue

* Type-stable `vect` implementation

Co-authored-by: Michael Abbott <[email protected]>

* type stable vect pullback

* Update src/rulesets/Base/array.jl

Co-authored-by: Michael Abbott <[email protected]>

* Test Union{Number, AbstractArray}

* Don't test inference below 1.6

* Update src/rulesets/Base/array.jl

Co-authored-by: Michael Abbott <[email protected]>

* Update src/rulesets/Base/array.jl

Co-authored-by: Michael Abbott <[email protected]>

* Style fix

* Style fix

* Update src/rulesets/Base/array.jl

Co-authored-by: Michael Abbott <[email protected]>

* Style fix

* Add an extra test

Co-authored-by: Michael Abbott <[email protected]>

Author: willtebbutt

src/rulesets/Base/array.jl

@@ -1,3 +1,42 @@
+#####
+##### constructors
+#####
+
+ChainRules.@non_differentiable (::Type{T} where {T<:Array})(::UndefInitializer, args...)
+
+function rrule(::Type{T}, x::AbstractArray) where {T<:Array}
+    project_x = ProjectTo(x)
+    Array_pullback(ȳ) = (NoTangent(), project_x(ȳ))
+    return T(x), Array_pullback
+end
+
+#####
+##### `vect`
+#####
+
+@non_differentiable Base.vect()
+
+# Case of uniform type `T`: the data passes straight through,
+# so no projection should be required.
+function rrule(::typeof(Base.vect), X::Vararg{T, N}) where {T, N}
+    vect_pullback(ȳ) = (NoTangent(), NTuple{N}(ȳ)...)
+    return Base.vect(X...), vect_pullback
+end
+
+# Numbers and arrays are often promoted, to make a uniform vector.
+# ProjectTo here reverses this
+function rrule(
+    ::typeof(Base.vect),
+    X::Vararg{Union{Number,AbstractArray{<:Number}}, N},
+) where {N}
+    projects = map(ProjectTo, X)
+    function vect_pullback(ȳ)
+        X̄ = ntuple(n -> projects[n](ȳ[n]), N)
+        return (NoTangent(), X̄...)
+    end
+    return Base.vect(X...), vect_pullback
+end
+
 #####
 ##### `reshape`
 #####

test/rulesets/Base/array.jl

@@ -1,3 +1,40 @@
+@testset "constructors" begin
+
+    # We can't use test_rrule here (as it's currently implemented) because the elements of
+    # the array have arbitrary values. The only thing we can do is ensure that we're getting
+    # `ZeroTangent`s back, and that the forwards pass produces the correct thing still.
+    # Issue: https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/202
+    @testset "undef" begin
+        val, pullback = rrule(Array{Float64}, undef, 5)
+        @test size(val) == (5, )
+        @test val isa Array{Float64, 1}
+        @test pullback(randn(5)) == (NoTangent(), NoTangent(), NoTangent())
+    end
+    @testset "from existing array" begin
+        test_rrule(Array, randn(2, 5))
+        test_rrule(Array, Diagonal(randn(5)))
+        test_rrule(Matrix, Diagonal(randn(5)))
+        test_rrule(Matrix, transpose(randn(4)))
+        test_rrule(Array{ComplexF64}, randn(3))
+    end
+end
+
+@testset "vect" begin
+    test_rrule(Base.vect)
+    @testset "homogeneous type" begin
+        test_rrule(Base.vect, (5.0, ), (4.0, ))
+        test_rrule(Base.vect, 5.0, 4.0, 3.0)
+        test_rrule(Base.vect, randn(2, 2), randn(3, 3))
+    end
+    @testset "inhomogeneous type" begin
+        test_rrule(
+            Base.vect, 5.0, 3f0;
+            atol=1e-6, rtol=1e-6, check_inferred=VERSION>=v"1.6",
+        ) # tolerance due to Float32.
+        test_rrule(Base.vect, 5.0, randn(3, 3); check_inferred=false)
+    end
+end
+
 @testset "reshape" begin
     test_rrule(reshape, rand(4, 5), (2, 10))
     test_rrule(reshape, rand(4, 5), 2, 10)