Re-enable EnzymeAdjoint testing

[ChrisRackauckas]

The direct adjoint of Enzyme is working again:

using Enzyme, OrdinaryDiffEq, StaticArrays

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]
end

const _saveat =  SA[0.0,0.25,0.5,0.75,1.0,1.25,1.5,1.75,2.0,2.25,2.5,2.75,3.0]

function f(y::Array{Float64}, u0::Array{Float64})
    tspan = (0.0, 3.0)
    prob = ODEProblem{true, SciMLBase.FullSpecialize}(lorenz!, u0, tspan)
    sol = DiffEqBase.solve(prob, Tsit5(), saveat = _saveat, sensealg = DiffEqBase.SensitivityADPassThrough())
    y .= sol[1,:]
    return nothing
end;
u0 = [1.0; 0.0; 0.0]
d_u0 = zeros(3)
y  = zeros(13)
dy = zeros(13)

Enzyme.autodiff(Reverse, f,  Duplicated(y, dy), Duplicated(u0, d_u0));

However, the EnzymeAdjoint which should be doing the same thing is hitting some odd error message: enzymerror.txt

[ChrisRackauckas]

@wsmoses can I get some help isolating this? Here's the MWE. Somehow autodiff works, but the autodiff_thunk call fails:

using Enzyme, OrdinaryDiffEq, StaticArrays

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]
end

const _saveat =  SA[0.0,0.25,0.5,0.75,1.0,1.25,1.5,1.75,2.0,2.25,2.5,2.75,3.0]

function f(y::Array{Float64}, u0::Array{Float64})
    tspan = (0.0, 3.0)
    prob = ODEProblem{true, SciMLBase.FullSpecialize}(lorenz!, u0, tspan)
    sol = DiffEqBase.solve(prob, Tsit5(), saveat = _saveat, sensealg = DiffEqBase.SensitivityADPassThrough())
    y .= sol[1,:]
    return nothing
end;
u0 = [1.0; 0.0; 0.0]
d_u0 = zeros(3)
y  = zeros(13)
dy = zeros(13)

Enzyme.autodiff(Reverse, f,  Duplicated(y, dy), Duplicated(u0, d_u0)); # This works!

f2 = (u0, p) -> solve(prob, Tsit5(); u0 = u0, p = p,
        sensealg = SensitivityADPassThrough(),
        saveat = _saveat)

splitmode = Enzyme.set_runtime_activity(Enzyme.ReverseSplitWithPrimal)

forward, reverse = Enzyme.autodiff_thunk(splitmode, Enzyme.Const{typeof(f2)}, Enzyme.Duplicated, Enzyme.Duplicated{typeof(u0)}, Enzyme.Duplicated{typeof(p)})
tape, result, shadow_result = forward(Enzyme.Const(f2), Enzyme.Duplicated(copy(u0), du0), Enzyme.Duplicated(copy(p), dp)) # Errors

[ChrisRackauckas]

enzymeerror.txt

[wsmoses]

can you get a new log with EnzymeAD/Enzyme.jl@03e3be8 I want to get some more info

[ChrisRackauckas]

enzymeerror2.txt

[wsmoses]

I tried to run the code above locally to inspect, got errors:

@ChrisRackauckas can you post a full MWE as an enzyme issue?

julia> forward, reverse = Enzyme.autodiff_thunk(splitmode, Enzyme.Const{typeof(f2)}, Enzyme.Duplicated, Enzyme.Duplicated{typeof(u0)}, Enzyme.Duplicated{typeof(p)})
ERROR: UndefVarError: `p` not defined
Stacktrace:
 [1] top-level scope
   @ REPL[12]:1

julia> tape, result, shadow_result = forward(Enzyme.Const(f2), Enzyme.Duplicated(copy(u0), du0), Enzyme.Duplicated(copy(p), dp)) # Errors
ERROR: UndefVarError: `du0` not defined
Stacktrace:
 [1] top-level scope
   @ REPL[13]:1

[ChrisRackauckas]

The main test here was changed to:

using SciMLSensitivity, OrdinaryDiffEq, Zygote, Enzyme

function fiip(du, u, p, t)
    du[1] = dx = p[1] * u[1] - p[2] * u[1] * u[2]
    du[2] = dy = -p[3] * u[2] + p[4] * u[1] * u[2]
end

p = [1.5, 1.0, 3.0, 1.0];
u0 = [1.0; 1.0];
prob = ODEProblem(fiip, u0, (0.0, 10.0), p)

du08, dp8 = Zygote.gradient(
    (u0, p) -> sum(solve(prob, Tsit5(), u0 = u0, p = p,
        abstol = 1e-14, reltol = 1e-14,
        saveat = 0.1,
        sensealg = EnzymeAdjoint())),
    u0,
    p)

No purely Enzyme MWE was found for what's going on here, those are all working now 😅.

[ChrisRackauckas]

Failure on v1.10 seems to be an Enzyme issue with the update that was supposed to fix this. v1.11 seems to be fine though.