rrule_via_ad // frule_via_ad Calling back into AD from ChainRules

[oxinabox]

This was originally discussed in JuliaDiff/ChainRules.jl#12 (comment)
and in a few other places.

Basically often when defining a chainrule (frule, or rrule)
it would be nice to be able to say "Give me this chainrule for some of the function, and if there is not one predefined us AD to get it"
as part of your rule definition.
Right now this is not ppossible, except by hard coding a AD (e.g. Zygote) in.

Where as if we had a function that an AD system could basically overload,
then we could do that.

It would also provide a common API for all ADs that support it.

This would help e.g. with higher order functions like map, broadcast etc
JuliaDiff/ChainRules.jl#122

There is some fiddlyness involved, around making sure it is both overloadable multiple times and so the user can choise which AD, and that compiles away, but I think we can sort it all out.

---

@jessebett reminded me of this today

[oxinabox]

Solution that @willtebbutt and I came up with the other night.
Make all rules take an extra argument for the configuration.
Which would be a struct:

struct Config{F, R}
    ad_frule::F
    ad_rrule::R
end

Where ad_frule and ad_rrule are functions conforming to the frule/rrule API, but that invoke a forward mode or reverse more AD.
Or they can be nothing if either one is not provided.

We might want to specify that they have to be actual functions so that we can actually dispatch on them being present.
e.g. for rrule

function rrule(::typeof(map), config::Config{<:Any, <:Function}, f, xs)
    y = map(f, xs...)
    function map_pullback(ȳ)
        ntuple(length(xs)+2) do full_i
            full_i == 1 && return NO_FIELDS
            full_i == 2 && return DoesNotExist()
            i = full_i-2
            @thunk map(ȳ, xs...) do ȳi, xis...
                _, pullback = ad_rrule(f, xis...)
                ∂xis = pullback(ȳi)
                extern(∂xis[i+1])  #+1 to skp ∂self
            end
        end
    end
    return y, map_pullback
end

Which since we dispatched on Confiig having a ad_rrule function it will hit the generic nothing rule-not-found fallback if it doesn't.
If it does, then we can assume the ad_rrule will itself check for actual rrules.

One worry is the case that one is in some non-AD senario, but where one knows all the things should have rules.
I think for that case the user can set the Config to use the current checked_rrule/ checked_ffrule which errors if the rule is not found.

[oxinabox]

@jrevels 's old idea that was in the source code for ages and i have ment to transfer is

In some weird ideal sense, the fallback for e.g. frule should actually be "get
the derivative via forward-mode AD". This is necessary to enable mixed-mode
rules, where e.g. frule is used within a rrule definition. For example,
broadcasted functions may not themselves be forward-mode primitives, but are
often forward-mode differentiable.
ChainRulesCore, by design, is decoupled from any specific AD implementation. How,
then, do we know which AD to fall back to when there isn't a primitive defined?
Well, if you're a greedy AD implementation, you can just overload frule and/or
rrule to use your AD directly. However, this won't play nice with other AD
packages doing the same thing, and thus could cause load-order-dependent
problems for downstream users.
It turns out, Cassette solves this problem nicely by allowing AD authors to
overload the fallbacks w.r.t. their own context. Example using ForwardDiff:

using ChainRulesCore, ForwardDiff, Cassette
Cassette.@context MyChainRuleCtx
# ForwardDiff, itself, can call `my_frule` instead of
# `frule` to utilize the ForwardDiff-injected ChainRulesCore
# infrastructure
my_frule(args...) = Cassette.recurse(MyChainRuleCtx(), frule, args...)
function Cassette.overdub(::MyChainRuleCtx, ::typeof(frule), f, x::Number)
    r = frule(f, x)
    if isa(r, Nothing)
        fx, df = (f(x), (_, Δx) -> ForwardDiff.derivative(f, x) * Δx)
    else
        fx, df = r
    end
    return fx, df
end

Which could work. It would basically fix any used of checked_rrule and checked_frule in existing code, but rewriting them not to ever error since frule would never return nothing.

I know for purposes of allowing 2nd derivatives etc
@YingboMa is already overdubbiing frule in ForwardDiff2
Which either makes this easier, because its just add an extra thing to existing overdub,
Or harder, because it ends up adding an extra layer of overdub, and bad things to performance happen when you nest Cassette.

[willtebbutt]

A slight variation on the Config option discussed above.

Firstly, add a couple of extra types:

abstract type AbstractAD{T} end

struct ForwardsModeAD{T} <: AbstractAD{T}
    pushforward::T
end

struct ReverseModeAD{T} <: AbstractAD{T}
    pullback::T
end

These simply wrap an AD e.g. Zygote.pullback to produce a thing that ChainRules knows about. ChainRules will then assume an API for the function.

Implement a fallback definition of frule and rrule:

frule(::AbstractAD, tangents, values...) = frule(tangents, values...)

rrule(::AbstractAD, values...) = rrule(values...)

This gives rule-authors two options.

1. implement a new method of frule or rrule that completely ignores any AD information.
2. implement a new method of frule or rrule that exploits whichever AD is passed in.

This gives AD package authors options various options:

• default to always passing in their AD
• default to not passing in any AD. They'll still get the AD-independent rules, but will have to rely on their AD being sufficiently clever to differentiate through stuff.
• any mixture of the above, and elect to use different ADs for particular functions e.g. the forwards-in-reverse trick for maping and broadcasting unary functions of a scalar.