Does Enzyme.jacobian support mutating operations as ForwardDiff.jacobian does?

[junyixu]

https://juliadiff.org/ForwardDiff.jl/v0.7/user/api.html#Jacobians-of-f(x::AbstractArray)::AbstractArray-1

An API like: ForwardDiff.jacobian(f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig = JacobianConfig(f!, y, x), check=Val{true}()) where where f! updates y.

[wsmoses]

No, but you can use the more general autodiff interface which handles arbitrary mutation/data structures.

[junyixu]

I give examples of forward and reverse modes here.

function f!(y::Vector{Float64}, x::Vector{Float64})
    y[1] = x[1] * x[1] + x[2] * x[1]
    y[2] = x[2] * x[2] + x[1] * x[2]
    return nothing
end

x  = [2.0, 3.0]
y  = [0.0, 0.0]

dx=onehot(x)
dy=onehot(y)

Enzyme.autodiff(Forward, f!, BatchDuplicated(y, dy), BatchDuplicated(x, dx))

x  = [2.0, 3.0]
y  = [0.0, 0.0]
dx=([0.0, 0.0], [0.0, 0.0])
dy=onehot(y)

Enzyme.autodiff(Reverse, f!, BatchDuplicated(y, dy), BatchDuplicated(x, dx))

But is there a plan for an API like ForwardDiff.jacobian?

[gdalle]

I think if you want the specific API of ForwardDiff.jl, i.e. functions x -> y represented as f!(y, x), you will be better served by https://github.com/gdalle/DifferentiationInterface.jl

The Enzyme bindings are probably not optimal yet (in particular, the Jacobian is not batched), but it will do what you want

[wsmoses]

@junyixu yeah we can add that if you prefer. I hadn't seen that API before so that's why we don't have it.

If you're interested in helping contribute, I'd be happy to help you with a PR.

[vchuravy]

@gdalle this should look like:

DifferentiationInterface.jacobian(f, y, AutoEnzyme(), x) 

right?

[gdalle]

Yes, and you can also specify the mode in AutoEnzyme, otherwise it will pick forward mode by default. It's not yet super optimized and it doesn't yet use batches but at least it's there

[wsmoses]

if this is sufficiently common, we should just add this API to Enzyme.

[gdalle]

I wouldnt mind, whenever the API exists upstream I try to use it. Today I added the reverse Jacobian and the HVP that you provide (https://github.com/gdalle/DifferentiationInterface.jl/pull/445)

[wsmoses]

@gdalle do you want to open a PR with your current implementation as a starting place?

[gdalle]

The issue is that my current implementation is not specific to Enzyme: it relies heavily on the DI machinery to derive jacobian from pushforward or pullback. The magic happens in this file but I don't think it's gonna be very heplful for a pure-Enzyme version.