SciML / SciML/DiffEqGPU.jl

Enzyme gradients through parameter-dependent continuous event times are incorrect

Open
#533 1 comment 0 reactions 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.