How to use AbstractDifferentiation as a user?

[scheidan]

Thanks for working on AbstractDifferentiation! It tackles a very relevant practical problem.

Unfortunately, I cannot figure out from the Readme how AbstractDifferentiation should be used.

The Readme says

"To use AbstractDifferentiation, first construct a backend instance ab::AD.AbstractBackend using your favorite differentiation package in Julia that supports AbstractDifferentiation."

but does not explain how to do this, see my failed attempt below (inspired from the test code).
Also,iIt would be nice if I could list all available backends somehow.

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)


# Is this the correct way to create a Backend?
struct ForwardDiffBackend1 <: AD.AbstractForwardMode end
const forwarddiff_backend1 = ForwardDiffBackend1()

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

# both fail with:
# MethodError: no method matching adjoint(::Nothing)
AD.gradient(zygote_backend1, foo, x) 
AD.gradient(forwarddiff_backend1, foo, x)

[frankschae]

We want AD package developers to depend on AbstractDifferentiation to define one primitive operation (and potentially the primal_value if needed). Currently, it is not implemented in any lower-level AD package yet. So, at the moment, you also need to add the primitive operation. That is missing in your attempt above.

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)

Also,iIt would be nice if I could list all available backends somehow.

Absolutely! The next steps include adding more primitives for different AD packages as test cases to check if all macros etc., in AbstractDifferentiation are fine. Then, AD package developers could go ahead and add these primitives to their package (as it is done above). Afterwards, end-users would only have to call

const zygote_backend = ZygoteBackend(),

where the primitives are then defined in Zygote.jl, and get all listed "derivative operations" like Jacobian from the API of AbstractDifferentiation.

[ChrisRackauckas]

Could the connections be added via Requires and then upstreamed?

[gdalle]

Does the current dev documentation solve this issue?