Making Enzyme.jl work with the integrator interface of OrdinaryDiffEq.jl

New to Julia
,
5

Julia Version 1.10.8
[7da242da] Enzyme v0.13.30
[1dea7af3] OrdinaryDiffEq v6.91.0
[1ed8b502] SciMLSensitivity v7.74.0

I also changed your code a little while trying to make it work on 1.11 (did not make it : issue SciMLSensitivity + Enzyme + 1.11 issue · Issue #2318 · EnzymeAD/Enzyme.jl · GitHub),

using OrdinaryDiffEq
using Enzyme
using SciMLSensitivity

function fun(du,u, p, t)
    du .= -p[1] * t
    nothing
end

function test_fun(p,prob)
    prob = remake(prob, p=p)
    sol = solve(prob,RK4(),save_everystep=false)
    res = sol.u[2]
    return res[1]
end


p = [1.0]
dp = [0.0]
u0 = [2.0]
prob = ODEProblem{true}(fun, u0, (1.0, 2.0), p)
dprob = Enzyme.make_zero(prob)

Enzyme.autodiff(Enzyme.Reverse, test_fun, Duplicated(p, dp),DuplicatedNoNeed(prob,dprob))
@info dp

The only thing that you can’t go around is making u0 a vector I think, because SciMLSensitivity needs to go through your du (must be a Ref), also always diff a remake not the original ODEProblem creation, it may work but its really not beautiful this way. You could go back to an explicit def for “fun” but since I made u0 a vector its better to avoid it

1 Like