add tests from FluxML#769

mcabbott · May 8, 2021 · 949d21f · 949d21f
1 parent bc78231
commit 949d21f
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/test/utils.jl b/test/utils.jl
@@ -70,4 +70,21 @@ end
   f2(x) = ForwardDiff.jacobian(x -> log.(x[1:3] .+ x[2:4]), x)
   x2 = rand(5) .+ 1
   @test Zygote.jacobian(f2, x2)[1] ≈ ForwardDiff.jacobian(f2, x2)
+
+  # Tests from https://github.com/FluxML/Zygote.jl/issues/769
+  f(x) = [2x[1]^2 + x[1],x[2]^2 * x[1]]
+  g1(x) = sum(ForwardDiff.jacobian(f,x))
+  out,back = Zygote.pullback(g1,[2.0,3.2])
+  stakehouse = back(1.0)[1]
+  @test typeof(stakehouse) <: Vector
+  @test size(stakehouse) == (2,)
+  @test stakehouse ≈ ForwardDiff.gradient(g1,[2.0,3.2])
+
+  g2(x) = prod(ForwardDiff.jacobian(f,x))
+  out,back = Zygote.pullback(g2,[2.0,3.2])
+  @test_skip back(1.0)[1] == ForwardDiff.gradient(g2,[2.0,3.2])  # contains NaN, @adjoint prod isn't careful
+
+  g3(x) = sum(abs2,ForwardDiff.jacobian(f,x))
+  out,back = Zygote.pullback(g3,[2.0,3.2])
+  @test back(1.0)[1] == ForwardDiff.gradient(g3,[2.0,3.2])
 end