Merge pull request #1039 from mzgubic/mz/cr1

Extra fixes for ChainRulesCore @1.0

Author: oxinabox

Project.toml

@@ -24,7 +24,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 [compat]
 AbstractFFTs = "0.5, 1.0"
 ChainRules = "0.8.12"
-ChainRulesCore = "0.10.4"
+ChainRulesCore = "1"
 ChainRulesTestUtils = "0.7.1"
 DiffRules = "1.0"
 FillArrays = "0.8, 0.9, 0.10, 0.11, 0.12"

src/compiler/chainrules.jl

@@ -111,7 +111,7 @@ for T_outer in (:Tuple, :NamedTuple)
   # than happy.
   @eval @inline function wrap_chainrules_output(x::ChainRules.Tangent{P, T}) where {P, T<:$T_outer}
     xp = map(wrap_chainrules_output, canonicalize(x))
-    convert($T_outer, xp)
+    ChainRulesCore.backing(xp)  # this is accessing ChainRulesCore internals, but it is prob safe enough, and it is fastest
   end
 end
 

test/complex.jl

@@ -18,7 +18,7 @@ using Zygote, Test, LinearAlgebra
 @test gradient(x -> real(logabsdet(x)[1]), [1 2im; 3im 4])[1] ≈ [4 3im; 2im 1]/10
 
 # https://github.com/FluxML/Zygote.jl/issues/705
-@test gradient(x -> imag(sum(exp, x)), [1,2,3])[1] ≈ im .* exp.(1:3)
+@test gradient(x -> imag(sum(exp, x)), [1,2,3])[1] ≈ real(im .* exp.(1:3))
 @test gradient(x -> imag(sum(exp, x)), [1+0im,2,3])[1] ≈ im .* exp.(1:3)
 
 fs_C_to_R = (real,

test/features.jl

@@ -449,12 +449,12 @@ end
 @test pullback(type_test)[1] == Complex{<:Real}
 
 @testset "Pairs" begin
-    @test (x->10*pairs((a=x, b=2))[1])'(100) === 10
+    @test (x->10*pairs((a=x, b=2))[1])'(100) === 10.0
     @test (x->10*pairs((a=x, b=2))[2])'(100) === 0
     foo(;kw...) = 1
     @test gradient(() -> foo(a=1,b=2.0)) === ()
 
-    @test (x->10*(x => 2)[1])'(100) === 10
+    @test (x->10*(x => 2)[1])'(100) === 10.0
     @test (x->10*(x => 2)[2])'(100) === 0
 end
 

test/gradcheck.jl

@@ -81,7 +81,7 @@ end
   @test gradient(xs ->sum(xs .^ _pow), [4, -1]) == ([_pow*4^9, -10],)
 
   @test gradient(x -> real((1+3im) * x^2), 5+7im) == (-32 - 44im,)
-  @test gradient(p -> real((1+3im) * (5+7im)^p), 2)[1] ≈ (-234 + 2im)*log(5 - 7im)
+  @test gradient(p -> real((1+3im) * (5+7im)^p), 2)[1] ≈ real((-234 + 2im)*log(5 - 7im))
   # D[(1+3I)x^p, p] /. {x->5+7I, p->2} // Conjugate
 end
 
@@ -160,7 +160,7 @@ end
 
   # https://github.com/FluxML/Zygote.jl/issues/376
   _, back = Zygote._pullback(x->x[1]*im, randn(2))
-  @test back(1.0)[2] == [-im, 0]
+  @test back(1.0)[2] == real([-im, 0]) == [0, 0]
 
   # _droplike
   @test gradient(x -> sum(inv, x[1, :]'), ones(2, 2)) == ([-1 -1; 0 0],)

test/utils.jl

@@ -24,10 +24,25 @@ end
 
   xs, y = randn(2,3), rand()
   f34(xs, y) = xs[1] * (sum(xs .^ (1:3)') + y^4)  # non-diagonal Hessian, two arguments
+
+  # Follow is should work ones we workout what ForwardDiff should do when `Float64` is called on a `Dual`
+  # https://github.com/JuliaDiff/ForwardDiff.jl/pull/538
+  # else might need a custom overload of `(;;ChainRulesCore.ProjectTo)(::Dual)`
+  # When fixed uncomment  the below and delete the broken function
+  #==
   dx, dy = diaghessian(f34, xs, y)
   @test size(dx) == size(xs)
   @test vec(dx) ≈ diag(hessian(x -> f34(x,y), xs))
   @test dy ≈ hessian(y -> f34(xs,y), y)
+  ==#  
+  function broken()
+    dx, dy = diaghessian(f34, xs, y) # This fails becase ProjectTo can't project a Dual onto a Float
+    c1 = size(dx) == size(xs)
+    c2 = vec(dx) ≈ diag(hessian(x -> f34(x,y), xs))
+    c3 = dy ≈ hessian(y -> f34(xs,y), y)
+    return all([c1, c2, c3])
+  end
+  @test_broken broken()
 
   zs = randn(7,13)  # test chunk mode
   @test length(zs) > ForwardDiff.DEFAULT_CHUNK_THRESHOLD