Crashes with `ComponentArrays.jl`, `DifferentialEquations.jl`, `SciMLSensitivity.jl` and `Enzyme.jl` on Julia 1.11.9

[GodotMisogi]

MWE:

# %%
using DifferentialEquations
using ComponentArrays
import DifferentiationInterface as DI
import SciMLSensitivity as SMS
import Enzyme
using Random

# Chaotic, so gradients are unreliable, but just for fun
function lorenz!(du, u, p, t)
    du[1] = p[1] * (u[2] - u[1])
    du[2] = u[1] * (p[2] - u[3]) - u[2]
    du[3] = u[1] * u[2] - p[3] * u[3]
    return nothing
end

# %%
const u0 = ComponentVector(x=1.0, y=0.0, z=0.0)
p_true = ComponentArray(sigma = 10.0, rho = 28.0, beta = 8.0/3.0)
p_init = ComponentArray(sigma = 5.0, rho = 15.0, beta = 1.0)
tspan = (0.0, 10.0)

prob_true = ODEProblem(lorenz!, u0, tspan, p_true)
sol_true = solve(prob_true, Tsit5(), saveat=0.1)

# %%
Random.seed!(123)
const data = Array(sol_true) + 2.0 * randn(size(Array(sol_true)))

# %%
function loss(p)
    sol = solve(prob_true, Tsit5(); p, saveat=0.1)
    return sum(abs2, Array(sol) - data) 
end
backend = SMS.AutoEnzyme() # Works
dp = DI.gradient(loss, backend, p_init) 

Results: Passing p_init instead of p_init[:] crashes the REPL.

[yolhan83]

this works :

# %%
using DifferentialEquations
using ComponentArrays
import DifferentiationInterface as DI
import SciMLSensitivity as SMS
import Enzyme
using Random

# Chaotic, so gradients are unreliable, but just for fun
function lorenz!(du, u, p, t)
    du[1] = p[1] * (u[2] - u[1])
    du[2] = u[1] * (p[2] - u[3]) - u[2]
    du[3] = u[1] * u[2] - p[3] * u[3]
    return nothing
end

# %%
const u0 = ComponentVector(x=1.0, y=0.0, z=0.0)
p_true = ComponentArray(sigma = 10.0, rho = 28.0, beta = 8.0/3.0)
p_init = ComponentArray(sigma = 5.0, rho = 15.0, beta = 1.0)
tspan = (0.0, 10.0)

prob_true = ODEProblem(lorenz!, u0, tspan, p_true)
sol_true = solve(prob_true, Tsit5(), saveat=0.1)

# %%
Random.seed!(123)
const data = Array(sol_true) + 2.0 * randn(size(Array(sol_true)))

# %%
function loss(p,prob_true,data)
    sol = solve(prob_true, Tsit5(); p, saveat=0.1)
    return sum(abs2, Array(sol) - data) 
end
backend = SMS.AutoEnzyme(;mode=Enzyme.set_runtime_activity(Enzyme.Reverse)) # Works
dp = DI.gradient(loss, backend, p_init, DI.Constant(prob_true),DI.Constant(data)) 

with Enzyme.jl I would avoid non-constant global args and use runtime activity because it may use Closure. If you're not used to write very type stable code, you may need to switch to Mooncake which can (sometimes) overcome those issues.

PS : when so much packages are acting together I would post on discouse https://discourse.julialang.org/latest before writing an issue to make sure where the problem comes from. So if this is ok close the issue and if you have any other question on your code don't hesitate to open a discourse thread.

[GodotMisogi]

Understood, thank you for the fix!