Merge 0f0d928 into c8679c6

Project.toml

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

src/rulesets/Base/array.jl

@@ -3,15 +3,17 @@
 #####
 
 function rrule(::typeof(reshape), A::AbstractArray, dims::Tuple{Vararg{Int}})
+    A_dims = size(A)
     function reshape_pullback(Ȳ)
-        return (NO_FIELDS, reshape(Ȳ, dims), DoesNotExist())
+        return (NO_FIELDS, reshape(Ȳ, A_dims), DoesNotExist())
     end
     return reshape(A, dims), reshape_pullback
 end
 
 function rrule(::typeof(reshape), A::AbstractArray, dims::Int...)
+    A_dims = size(A)
     function reshape_pullback(Ȳ)
-        ∂A = reshape(Ȳ, dims)
+        ∂A = reshape(Ȳ, A_dims)
         return (NO_FIELDS, ∂A, fill(DoesNotExist(), length(dims))...)
     end
     return reshape(A, dims...), reshape_pullback

test/rulesets/Base/array.jl

@@ -1,23 +1,11 @@
 @testset "reshape" begin
-    A = rand(4, 5)
-    B, pullback = rrule(reshape, A, (5, 4))
-    @test B == reshape(A, (5, 4))
-    Ȳ = randn(4, 5)
+    x = rand(4, 5)
+    x̄ = rand(4, 5)
 
-    (s̄, Ā, d̄) = pullback(Ȳ)
-    @test s̄ == NO_FIELDS
-    @test d̄ isa DoesNotExist
-    @test extern(Ā) == reshape(Ȳ, (5, 4))
+    ȳ = rand(2, 10)
 
-    B, pullback = rrule(reshape, A, 5, 4)
-    @test B == reshape(A, 5, 4)
-
-    Ȳ = randn(4, 5)
-    (s̄, Ā, d̄1, d̄2) = pullback(Ȳ)
-    @test s̄ == NO_FIELDS
-    @test d̄1 isa DoesNotExist
-    @test d̄2 isa DoesNotExist
-    @test extern(Ā) == reshape(Ȳ, 5, 4)
+    rrule_test(reshape, ȳ, (x, x̄), ((2, 10), nothing))
+    rrule_test(reshape, ȳ, (x, x̄), (2, nothing), (10, nothing))
 end
 
 @testset "hcat" begin