Support new ChainRulesCore.NoTangent

Author: oschulz

src/chain_rules_aliases.jl

@@ -12,3 +12,9 @@ else
     const Tangent{P,T} = ChainRulesCore.Composite{P,T}
     const ZeroTangent = ChainRulesCore.Zero
 end
+
+@static if isdefined(ChainRulesCore, :NoTangent)
+    const NoTangent = ChainRulesCore.NoTangent
+else
+    const NoTangent = ZeroTangent
+end

src/rrules.jl

@@ -18,7 +18,7 @@ end
 
 
 Base.@generated function _forwarddiff_pullback_thunks(f::Base.Callable, xs::NTuple{N,Any}, ΔΩ::Any) where N
-    Expr(:tuple, ChainRulesCore.NO_FIELDS, (:(ForwardDiffPullbacks.FwdDiffPullbackThunk(f, xs, Val($i), ΔΩ)) for i in 1:N)...)
+    Expr(:tuple, NoTangent(), (:(ForwardDiffPullbacks.FwdDiffPullbackThunk(f, xs, Val($i), ΔΩ)) for i in 1:N)...)
 end
 
 g_ΔΩ = nothing
@@ -56,7 +56,7 @@ end
 
 
 Base.@generated function _forwarddiff_bc_pullback_thunks(f::Base.Callable, Xs::NTuple{N,Any}, ΔΩA::Any) where N
-    Expr(:tuple, ChainRulesCore.NO_FIELDS, ChainRulesCore.NO_FIELDS, (:(ForwardDiffPullbacks.FwdDiffBCPullbackThunk(f, Xs, Val($i), ΔΩA)) for i in 1:N)...)
+    Expr(:tuple, NoTangent(), NoTangent(), (:(ForwardDiffPullbacks.FwdDiffBCPullbackThunk(f, Xs, Val($i), ΔΩA)) for i in 1:N)...)
 end
 
 function ChainRulesCore.rrule(::typeof(Base.broadcasted), wrapped_f::WithForwardDiff, Xs::Vararg{Any,N}) where N

src/with_forwarddiff.jl

@@ -32,7 +32,7 @@ using ChainRulesCore
 
 y, back = rrule(fwddiff(f), xs...)
 y == 139
-map(unthunk, back(1)) == (ZeroTangent(), 4, (6, 8), [10, 12, 14])
+map(unthunk, back(1)) == (NoTangent(), 4, (6, 8), [10, 12, 14])
 
 using Zygote
 

test/test_rrules.jl

@@ -9,6 +9,6 @@ include("testfuncs.jl")
 
 @testset "ChainRulesCore" begin
     @test @tinferred(ChainRulesCore.rrule(fwddiff(f), xs...)) isa Tuple{Tuple, Function}
-    @test @tinferred((ChainRulesCore.rrule(fwddiff(f), xs...)[2])(ΔΩ)) isa Tuple{ZeroTangent, ForwardDiffPullbacks.FwdDiffPullbackThunk, ForwardDiffPullbacks.FwdDiffPullbackThunk,ForwardDiffPullbacks.FwdDiffPullbackThunk}
-    @test @tinferred(map(unthunk, (ChainRulesCore.rrule(fwddiff(f), xs...)[2])(ΔΩ))) == (ZeroTangent(), 280, (600, 1040), SVector(1600, 2280, 3080))
+    @test @tinferred((ChainRulesCore.rrule(fwddiff(f), xs...)[2])(ΔΩ)) isa Tuple{NoTangent, ForwardDiffPullbacks.FwdDiffPullbackThunk, ForwardDiffPullbacks.FwdDiffPullbackThunk,ForwardDiffPullbacks.FwdDiffPullbackThunk}
+    @test @tinferred(map(unthunk, (ChainRulesCore.rrule(fwddiff(f), xs...)[2])(ΔΩ))) == (NoTangent(), 280, (600, 1040), SVector(1600, 2280, 3080))
 end

test/test_zygote.jl

@@ -13,7 +13,7 @@ include("testfuncs.jl")
         y, back = ChainRulesCore.rrule(f, xs...)
         back_thunks = back(ΔΩ)
         Δx = map(unthunk, back_thunks)
-        @assert first(Δx) == ZeroTangent()
+        @assert first(Δx) == NoTangent()
         y, Base.tail(Δx)
     end
 
@@ -28,7 +28,7 @@ include("testfuncs.jl")
         y, back = ChainRulesCore.rrule(Base.broadcasted, f, Xs...)
         back_thunks = back(ΔΩA)
         Δx = map(unthunk, back_thunks)
-        @assert first(Δx) == ZeroTangent()
+        @assert first(Δx) == NoTangent()
         y, Base.tail(Δx)
     end