Move forwarddiff and ignore to ChainRulesCore.jl?

[oschulz]

Would allow people to use hints like forwarddiff and ignore more liberally in their code, without adding Zygote itself as a dependency. I opened the same issue on ZygoteRules.jl, since it affects both packages: FluxML/ZygoteRules.jl#9

[oxinabox]

ignore should probably be moved to ChainRulesCore.jl.

[oschulz]

I have to profess my ignorance here - what's the relationship between ZygoteRules.jl and ChainRulesCore.jl? From the docs, it looks like both are intended for users to want to add AD hints to their code or implement custom gradients for their functions? Will ChainRulesCore replace ZygoteRules eventually?

[sethaxen]

ZygoteRules rules are unique to Zygote, while ChainRulesCore rules can in principle be used by any AD (currently Zygote and ForwardDiff2 use them, but I believe others are planned).

As for when you should use one vs the other, prefer to use ChainRulesCore unless your rule relies on some Zygote-specific feature like calling Zygote.pullback within the rule (not currently supported by ChainRulesCore) or using Zygote.Buffer for a mutating function. Any (reverse-mode for now) rule defined with ChainRulesCore will be used by Zygote as well.

[oschulz]

Thanks for the clarification!

Then I guess we'd want to have both the forwarddiff and the ignore hint in ChainRulesCore? ChainRulesCore presents forward and reverse mode as concepts, so I think both hints would fit into it's scope.

I believe people would definitely be encouraged to add more AD-hinting to their code if that only requires depending on something as lightweight and generic as ChainRulesCore.

[oschulz]

@oxinabox what's your opinion on this? I recently found myself in a situation where I need to add Zygote.forwarddiff hinting to make code Zygote-compatible, but I don't want to add Zygote as a dep just because of this if avoidable. Could forwarddiff() live in ZygoteRules? Or could there be some generic "switch to forward mode" hinting capability in ChainRulesCore?

[oxinabox]

AFAIK Zygote.forwarddiff basically just calls ForwardDiff.jl.
So you can probably comble something together using that directly.
(This also has implications on where it can live since things don't want a ForwardDiff.jl dependency)

Could forwarddiff() live in ZygoteRules?

I won't comment on on this since i don't maintain ZygoteRules, and also strongly recommend wherever possible people stop using it. (using ChainRulesCore instead).

Or could there be some generic "switch to forward mode" hinting capability in ChainRulesCore?

This is something that will become possible once we have
JuliaDiff/ChainRulesCore.jl#68
which will let us specify inside an rrule instructions to compute some part of the problem (or all of the problem) using the forward-mode AD.
This issue is one of the last really big things ChainRulesCore needs to solve.

[oschulz]

AFAIK Zygote.forwarddiff basically just calls ForwardDiff.jl.

Just curious - this may change at some point though? Zygote is getting it's own forward-mode diff, I heard?

This also has implications on where it can live since things don't want a ForwardDiff.jl dependency

Ah, right, we certainly wouldn't want a ForwardDiff dependency in there.

This is something that will be come possible once we have JuliaDiff/ChainRulesCore.jl#68

Sounds good, thanks! I was asking more for the long term anyhow.