Using Enzyme with DE.jl

I am playing with using Enzyme.jl to try to get the derivative of some numerical function with respect to the initial condition. The function is solved using DifferentialEquations.jl. For example:

using DifferentialEquations
using Enzyme 

"""
The differential equation 
"""
function lorenz!(du,u,p,t)
    du[1] = 10.0*(u[2]-u[1])
    du[2] = u[1]*(28.0-u[3]) - u[2]
    du[3] = u[1]*u[2] - (8/3)*u[3]
    nothing
  end

"""
Given initial conditions a,b,c, integrate the system lorenz!()
"""
function timestepping(a,b,c)

    u0 = [a,b,c] 
    tspan = (0.0,100.0)
   
    prob = ODEProblem(lorenz!,u0,tspan)
    
    ode_solution = solve(prob,Tsit5())

    return ode_solution

end 

"""
An arbitrary loss function that returns the final value of the `a` variable at the end of the integration
"""
function loss_function(a,b,c)

    #Do an integration 
    solution = timestepping(a,b,c)

    #Get just the last element of the first variable as an arbitrary loss 
    critical_value = last(solution[1,:])
    
    return critical_value

end 


a = 1.0
b = 2.0
c = 3.0

#Just run once, make sure everything works OK
first_attempt = loss_function(a,b,c)

#Now try to get a derivative
second_attempt = Enzyme.autodiff(Forward,loss_function,Duplicated,Duplicated(a,1.0),Const(b),Const(c))

i.e. I am trying to get the derivative with respect to the initial condition a, holding all other variables constant.

However this throws a LoadError: Tuple{DataType} is not a dispatch tuple which seems to be related to the GPUCompiler.specialization_id(job) call. Can anyone lend some guidance on where I am going wrong here? Thanks!

You’d want to do Enzyme just on the inside, i.e. EnzymeVJP, unless you’re on Enzyme master (which I assume you’re probably not :sweat_smile:)

Sorry Chris could you expand a little more? What exactly do you mean by “on the inside”?

If you just use Zygote on that it would use Enzyme for the vjps internally. Details in Sensitivity Algorithms for Differential Equations with Automatic Differentiation (AD) · SciMLSensitivity.jl. Enzyme should do this soon but needs the new rules system release.