Merge pull request #785 from JuliaDiff/ox/izt

support arbitary types for vect -- use zero_tangent for _instantiate_zeros

Author: oxinabox

src/rulesets/Base/array.jl

@@ -31,7 +31,7 @@ end
 
 @non_differentiable Base.vect()
 
-function frule((_, ẋs...), ::typeof(Base.vect), xs::Number...)
+function frule((_, ẋs...), ::typeof(Base.vect), xs...)
     return Base.vect(xs...), Base.vect(_instantiate_zeros(ẋs, xs)...)
 end
 
@@ -71,9 +71,7 @@ materialises each zero `ẋ` to be `zero(x)`.
 """
 _instantiate_zeros(ẋs, xs) = map(_i_zero, ẋs, xs)
 _i_zero(ẋ, x) = ẋ
-_i_zero(ẋ::AbstractZero, x) = zero(x)
-# Possibly this won't work for partly non-diff arrays, something like `gradient(x -> ["abc", x][end], 1)`
-# may give a MethodError for `zero` but won't be wrong.
+_i_zero(ẋ::AbstractZero, x) = zero_tangent(x)
 
 # Fast paths. Should it also collapse all-Zero cases?
 _instantiate_zeros(ẋs::Tuple{Vararg{Number}}, xs) = ẋs

test/rulesets/Base/array.jl

@@ -35,9 +35,16 @@ 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,))
+        test_frule(Base.vect, (5.0,), (4.0,))
         test_rrule(Base.vect, 5.0, 4.0, 3.0)
+        test_frule(Base.vect, 5.0, 4.0, 3.0)
         test_rrule(Base.vect, randn(2, 2), randn(3, 3))
+        test_frule(Base.vect, randn(2, 2), randn(3, 3))
+
+        # Nonnumber types
+        test_frule(Base.vect, (1.0, 2.0), (1.0, 2.0))
+        test_rrule(Base.vect, (1.0, 2.0), (1.0, 2.0))
     end
     @testset "inhomogeneous type" begin
         # fwd
@@ -52,7 +59,7 @@ end
     end
     @testset "_instantiate_zeros" begin
         # This is an internal function also used for `cat` etc.
-        @eval using ChainRules: _instantiate_zeros
+        _instantiate_zeros = ChainRules._instantiate_zeros
         # Check these hit the fast path, unrealistic input so that map would fail:
         @test _instantiate_zeros((true, 2 , 3.0), ()) == (1, 2, 3)
         @test _instantiate_zeros((1:2, [3, 4]), ()) == (1:2, 3:4)