Open junyixu opened 5 months ago
No, but you can use the more general autodiff interface which handles arbitrary mutation/data structures.
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
?
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
@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.
@gdalle this should look like:
DifferentiationInterface.jacobian(f, y, AutoEnzyme(), x)
right?
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
if this is sufficiently common, we should just add this API to Enzyme.
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)
@gdalle do you want to open a PR with your current implementation as a starting place?
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 wheref!
updatesy
.