Merge pull request #1205 from devmotion/dw/notimplemented

Handle `ChainRulesCore.NotImplemented`

Author: CarloLucibello

Project.toml

@@ -1,6 +1,6 @@
 name = "Zygote"
 uuid = "e88e6eb3-aa80-5325-afca-941959d7151f"
-version = "0.6.37"
+version = "0.6.38"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/compiler/chainrules.jl

@@ -106,6 +106,7 @@ Convert `x` from the differentials types ChainRules uses to the format Zygote us
 # Zygote convention: even if many AbstractZero partials (i.e. multi-input function), make just 1 nothing.
 @inline wrap_chainrules_output(x::Tuple{Vararg{ChainRules.AbstractZero}}) = nothing
 @inline wrap_chainrules_output(x::ChainRules.AbstractZero) = nothing
+@inline wrap_chainrules_output(x::ChainRulesCore.NotImplemented) = nothing
 for T_outer in (:Tuple, :NamedTuple)
   # we create separate methods rather than using a `Union` + an `if` so that we avoid a
   # branch that changes output type, because nested AD on that kinda thing makes Zygote less

test/chainrules.jl

@@ -263,6 +263,18 @@ using Zygote: ZygoteRuleConfig
         @test (1.0,) == Zygote.gradient(oout_id_outer, π)
         @test oout_id_rrule_hitcount[] == 0
     end
+
+    # issue #1204
+    @testset "NotImplemented" begin
+        f_notimplemented(x) = x
+        @scalar_rule f_notimplemented(x) @not_implemented("not implemented :(")
+        @test Zygote.gradient(f_notimplemented, 0.1) === (nothing,)
+        @test Zygote.gradient(x -> f_notimplemented(x[1]), 0.1) === (nothing,)
+        if isdefined(Base, :only)
+            @test Zygote.gradient(x -> f_notimplemented(only(x)), (0.1,)) === (nothing,)
+            @test Zygote.gradient(x -> f_notimplemented(only(x)), [0.1]) === (nothing,)
+        end
+    end
 end
 
 @testset "ChainRulesCore.rrule_via_ad" begin