Add ChainRulesCore RuleConfig?

[gdalle]

The macro @grad_from_chainrules contains the following line

ReverseDiff.jl/src/macros.jl

Line 338 in 6c7b9cf

output_value, back = ChainRulesCore.rrule($f, map(ReverseDiff.value, args)...; $kwargs...)

which shows that the rrule is used with no optional RuleConfig argument. This is proving cumbersome for ImplicitDifferentiation.jl, where the rrule needs to call back into AD and thus dispatch on the AD backend (aka RuleConfig).

Would you be open to a PR defining the following config?

using ChainRulesCore: RuleConfig, HasReverseMode, NoForwardsMode

struct MyReverseDiffRuleConfig <: RuleConfig{Union{HasReverseMode,NoForwardsMode}} end

I would need advice on how to add it to the macro though.

Pinging @mohamed82008

[mohamed82008]

We would need a method for rrule_via_ad and the ReverseDiffRuleConfig first, similar to https://github.com/FluxML/Zygote.jl/blob/master/src/compiler/chainrules.jl#L248. Then we can change that line you pointed to.