How to use AbstractDifferentiation as a user? about abstractdifferentiation.jl HOT 2 OPEN

using AbstractDifferentiation
import Zygote
import ForwardDiff

# test function
foo(x) = sin(x[1]) + prod(x[2:end].^2)
x = rand(4)

# direct usage works
Zygote.gradient(foo, x)[1]
ForwardDiff.gradient(foo, x)

# correct way to use a backend
struct ForwardDiffBackend1 <: AD.AbstractForwardMode end
const forwarddiff_backend1 = ForwardDiffBackend1()

AD.@primitive function pushforward_function(ab::ForwardDiffBackend1, f, xs...)
    # jvp = f'(x)*v, i.e., differentiate f(x + h*v) wrt h at 0
    return function (vs)
        if xs isa Tuple
            @assert length(xs) <= 2
            if length(xs) == 1
                (ForwardDiff.derivative(h->f(xs[1]+h*vs[1]),0),)
            else
                ForwardDiff.derivative(h->f(xs[1]+h*vs[1], xs[2]+h*vs[2]),0)
            end
        else
            ForwardDiff.derivative(h->f(xs+h*vs),0)
        end
    end
end
AD.primal_value(::ForwardDiffBackend1, ::Any, f, xs) = ForwardDiff.value.(f(xs...))

struct ZygoteBackend1 <: AD.AbstractReverseMode end
const zygote_backend1 = ZygoteBackend1()

AD.@primitive function pullback_function(ab::ZygoteBackend1, f, xs...)
    return function (vs)
        # Supports only single output
        _, back = Zygote.pullback(f, xs...)
        if vs isa AbstractVector
            back(vs)
        else
            @assert length(vs) == 1
            back(vs[1])
        end
    end
end

AD.gradient(zygote_backend1, foo, x) 
AD.gradient(forwarddiff_backend1, foo, x)

const zygote_backend = ZygoteBackend(),