Use ChainRules RuleConfig

[oxinabox]

This PR's commit history is pretty gross.
built on top of the merge of #987 and #967
and then hacks them up.

Right now its kinda just a sketch.
It might compile and work, but I assume it doesn't.

This needs: JuliaDiff/ChainRules.jl#441 which it uses as a test-case.

[oxinabox]

Still trying to debug why it won't pick up new rrules til a Zygote.refresh() is done.
Generated functions in general don't see things defined after them,
but we should have the right extra back-edged in place to cause them to see them.
Zygote.refresh() fixes it, but we shouldn't need that, we didn't need it for this case before.

The following is what happens right now:
afterr defining the rrule for cr_inner_demo.

julia> Zygote._pullback(Zygote.Context(), cr_outer_demo, 11)[2](1)
T=Tuple{Zygote.ZygoteRuleConfig{Zygote.Context}, typeof(Main.cr_outer_demo), Int64}; return_type=Nothing
instance=rrule(Zygote.ZygoteRuleConfig{Zygote.Context}, typeof(Main.cr_outer_demo), Int64) from rrule(ChainRulesCore.RuleConfig{T} where T, Any, Any...)
T=Tuple{Zygote.ZygoteRuleConfig{Zygote.Context}, typeof(Main.cr_inner_demo), Int64}; return_type=Nothing
instance=rrule(Zygote.ZygoteRuleConfig{Zygote.Context}, typeof(Main.cr_inner_demo), Int64) from rrule(ChainRulesCore.RuleConfig{T} where T, Any, Any...)
T=Tuple{Zygote.ZygoteRuleConfig{Zygote.Context}, typeof(Base.:(*)), Int64, Int64}; return_type=Tuple{Int64, ChainRules.var"#times_pullback#1490"{Int64, Int64}}
T=Tuple{Zygote.ZygoteRuleConfig{Zygote.Context}, typeof(Base.:(*)), Int64, Int64}; return_type=Tuple{Int64, ChainRules.var"#times_pullback#1490"{Int64, Int64}}
(nothing, 50)

julia> Zygote._pullback(Zygote.Context(), cr_outer_demo, 11)[2](1)  # just uses what it found before by ADing fully wihtout rule
(nothing, 50)

The problem is this line

T=Tuple{Zygote.ZygoteRuleConfig{Zygote.Context}, typeof(Main.cr_inner_demo), Int64}; return_type=Nothing

but it should be (shown after doing a Zygote.refresh()

T=Tuple{Zygote.ZygoteRuleConfig{Zygote.Context}, typeof(Main.cr_inner_demo), Int64}; return_type=Tuple{Int64, Main.var"#cr_inner_demo_pullback#10"}

[oxinabox]

Alternative solution for detecting if rules don't exist, since we can't ose Core.Compiler.return_type.

First check if we are hitting the rrule(::RuleConfig, args...) = rrule(args...) redispatcher.
then check if the rrule(args...) would hit the rrule(args...) = nothing fallback.
If that is the case then we have no rule and so need to attach the backedge, to the redispatcher (which will be invalidated in turn by backedged from the fallback anyway, so no need to attach backedges there).

Otherwise, we have hit a rule and we need no edges (because we will be returning an Expr that calls the rule, so natural edges apply).

[oxinabox]

The core of this now works. All the tests in test/chainrules.jl pass.
Mostly what is left is chasing up edge cases in JuliaDiff/ChainRules.jl#441
of which there are surprisingly many (and getting that merged)

[oxinabox]

This is all passing for me locally, so just need JuliaDiff/ChainRules.jl#441 merged and it will pass on CI.

@mzgubic will review tomorrow with aim being to merge it tomorrow

[mzgubic]

Thanks, looks great, excited about this!

One remaining concern is how do we opt out of using rules if we can't rely on the return value? I guess one way is to call f/rrule_via_ad inside the rule that we are opting out from. But that is far from ideal as it flops between the rule and AD and this results in redundant differential type changes. Is there a better way?

[oxinabox]

One remaining concern is how do we opt out of using rules if we can't rely on the return value?

Correct, that concern remains.
It is out of scope for this PR to solve, i had hoped we could solve it incidentally, but it seems the obvious way doens't work.

I guess one way is to call f/rrule_via_ad inside the rule that we are opting out from.

This is not trivial since, it will (if nothing special is done) just infinite-loop.
Maybe invoke can be used to prevent that?