Enzyme gradients through parameter-dependent continuous event times are incorrect
Nobody has claimed this yet.
- Dominant language
- Julia
- Stars
- 327
- Forks
- 42
- Avg merge
- 14h 13m
- Merged PRs (30d)
- 34
Description
A continuous callback locates the event correctly in the forward solve, but Enzyme reverse mode returns incorrect event-time sensitivities. For u′=1, u(0)=0, and a reset to zero when u=p, with 0.5<p<1, exactly one event occurs before t=1. Thus u(1)=1-p and every parameter derivative must be -1.
Reproduced on Julia 1.12.7, Enzyme 0.13.201, KernelAbstractions 0.9.42, and SciMLBase 3.53.1, using the CPU backend and the Enzyme restoration branch from https://github.com/SciML/DiffEqGPU.jl/pull/532. This is a remaining limitation of that restoration, not a claim that its changes introduced the event-time problem. The unmodified base cannot execute the ensemble-gradient reproducer because it fails earlier in Enzyme compilation.
using DiffEqGPU, Enzyme, KernelAbstractions, SciMLBase, StaticArrays, Test
function continuous_loss(p)
prob = ODEProblem{false}((u, p, t) -> SVector(1.0), SVector(0.0), (0.0, 1.0), SVector(p[1]))
pf = (prob, ctx) -> remake(prob; p = SVector(p[ctx.sim_id]))
condition(u, t, integrator) = u[1] - integrator.p[1]
affect!(integrator) = (integrator.u = zero(integrator.u))
cb = ContinuousCallback(condition, affect!; save_positions = (false, false))
sol = solve(EnsembleProblem(prob; prob_func = pf, safetycopy = false), GPUTsit5(), EnsembleGPUKernel(CPU(), 0.0); trajectories = length(p), adaptive = false, dt = 0.05, callback = cb, merge_callbacks = true, save_everystep = false)
return sum(s -> sum(s.u[end]), sol.u)
end
p = [0.63, 0.71, 0.84]
@test continuous_loss(p) ≈ 3 - sum(p)
dp = zero(p)
Enzyme.autodiff(Reverse, continuous_loss, Active, Duplicated(p, dp))
@show dp
@test dp ≈ -ones(3) rtol = 1.0e-6
Run with julia +1.12 --project=<environment developed against the Enzyme branch> continuous_event_reproducer.jl.
The forward-value assertion passes. The derivative assertion fails:
dp = [-1.0000000000000027, -0.4999999999999896, -0.7599918704695224]
Test Failed:
Expression: dp ≈ -ones(3) rtol = 1e-6
Evaluated: [-1.0000000000000027, -0.4999999999999896, -0.7599918704695224] ≈ [-1.0, -1.0, -1.0] (rtol=1.0e-6)
The same probe with a fixed-time discrete callback (doubling the exponential solution at t=0.5, using tstops=[0.5] and merge_callbacks=true) passes its analytic gradient check. This distinction should be reflected in documentation: callback execution is supported, but that does not establish correct sensitivities through parameter-dependent event times. Root localization and the differentiated event update need investigation; no cause has been bisected, and GPU behavior has not been verified.
🤖 Investigated with Codex CLI 0.153.4 (model: gpt-6-astra; local session ID: 01a07fcc-1c4f-7ee3-9a1e-51eaab7df293).
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with continuous_event_reproducer.jl and the continuous_loss entry point, then inspect the ContinuousCallback and EnsembleGPUKernel(CPU()) path used by the reproducer. Run it with the stated Julia project and Enzyme branch, compare the continuous-event result with the passing fixed-time callback case, and investigate until the parameter derivatives are all -1.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- hpc, performance
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100