search

JuliaDiff / ReverseDiff.jl

Reverse Mode Automatic Differentiation for Julia
Other
348 stars 57 forks source link

Add ChainRulesCore RuleConfig? #238

Open gdalle opened 1 year ago

gdalle commented 1 year ago

The macro @grad_from_chainrules contains the following line

https://github.com/JuliaDiff/ReverseDiff.jl/blob/6c7b9cfbb06cc1e91d36be12746aaeb55e433ff4/src/macros.jl#L338

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 commented 1 year ago

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.