Zygote.gradient returns `nothing` instead of `NotImplemented

[simonmandlik]

using ChainRulesCore, Zygote

f(x) = 2x
ChainRulesCore.rrule(::typeof(f), x) = 2x, d -> (NoTangent(), @not_implemented(""), )

Zygote.gradient(f, 1)

The last line returns (nothing, ) from v0.6.38 onwards but used to return ChainRulesCore.NotImplemented in the previous versions.

Is this intended?

[ToucheSir]

Yes: #1205. As an alternative, you could try using rrule_via_ad(ZygoteRuleConfig(), f, 1). That won't help if a NotImplemented is generated in the process of running AD, but it will short-circuit if it finds a rrule.

If anyone has any ideas on how to refactor the internals such that rrule_via_ad doesn't have to go through an unwrapping/re-wrapping step (i.e. AD uses ChainRules types all the way through), let me know and I'd be happy to discuss + assist in implementation.

[simonmandlik]

Ok, thanks!

[ToucheSir]

Re-opening as a tracker for JuliaDiff/ChainRules.jl#521 (comment).