Add rrule_from_frule macro

[niklasschmitz]

This draft PR implements a proposal for a new rrule definition helper macro @rrule_from_frule. It generates an rrule from an already defined frule by calling-back into (reverse-mode) AD.

This is also inspired by JAX's approach to decomposing reverse-mode into (forward-mode) linearization + transposition (where transposition only needs rules for linear primitives).

This minimum working example illustrates the idea from a ChainRules.jl perspective:

# rrule from frule (transposition)
using Zygote
using ChainRulesCore
using LinearAlgebra

function f(x)
    a = sin.(x)
    b = sum(a)
    c = b * a
    return c
end

function ChainRulesCore.frule((Δself, Δx,), ::typeof(f), x)
    a, ȧ = sin.(x), cos.(x) .* Δx
    b, ḃ = sum(a), sum(ȧ)
    c, ċ = b * a, ḃ * a + b * ȧ
    return c, ċ
end

function ChainRulesCore.rrule(config::RuleConfig{>:HasReverseMode}, ::typeof(f), x)
    pushforward(Δfx...) = frule(Δfx, f, x)[2]
    _, back = rrule_via_ad(config, pushforward, f, x)
    f_pullback(Δy) = back(Δy)[2:end]
    return f(x), f_pullback
end

let x = rand(3)
    v = randn(3)
    w = randn(3)
    jvp(f, x, v) = frule((NoTangent(), v), f, x)[2]
    vjp(f, x, w) = rrule_via_ad(Zygote.ZygoteRuleConfig(), f, x)[2](w)[2]
    dot(w, jvp(f, x, v)) ≈ dot(vjp(f, x, w), v)
end

This PR aims to provide this approach as a macro that can be used like @rrule_from_frule f(x::AbstractArray{<:Real}) to minimize boiler-plate code.

[oxinabox]

Nice.

If we resolved JuliaDiff/ChainRulesTestUtils.jl#221
we could drop the Zygote dependency from testing which would be good, but not blocking for this.

It might be nice to add a short explination of why this works into the docstring.

My notes from slack are

The pushforward at the point is a linear function which has the same derivative as the primal at that point.
So asking for it's derivative gives you the derivative of the primal.
And it is probably easier to AD than the primal function since you have already done the work of removing nonlinear parts. (at least with aggressive constant folding they will be gone, idk that we have that)

[niklasschmitz]

Thanks @oxinabox! I now added a short docstring. I also have left two more TODO's in the code:

1. Allowing more general function signature expressions in the macro, importantly multiple arguments (and possibly varargs).
   Should the handling of these be modelled after the macros in ChainRulesCore.jl or would you recommend a package like ExprTools.jl or MLStyle.jl? This is my first macro so I'm very open for guiding suggestions.

2. Avoiding redundant computation of the primal. Currently, the frule is executed (inside the pushforward, discarding the primal value), and then additionally the primal is executed once to get the correct primal value to return from the rrule. I haven't solved it yet, but I think there's likely room for improvement here.

[oxinabox]

2. Avoiding redundant computation of the primal. Currently, the frule is executed (inside the pushforward, discarding the primal value), and then additionally the primal is executed once to get the correct primal value to return from the rrule. I haven't solved it yet, but I think there's likely room for improvement here.

Related to this, I think Zygote (and probably Diffractor) will AFAIK also generate code for the pullback of that discarded primal i think.
Unlike Jax we can't aggressively constant fold that out of the pushforward, though I haven't seen the new tools that are in Julia 1.8 (and maybe 1.7-rc2?) (cc @simeonschaub might know).