Frequent JIT-recompile of discrete_terminating_event

Hi,

I am running into a strange issue when using the diffrax.diffeqsolve with the discrete_terminating_event argument which I believe is due to a large number of JIT-recompiles, making execution time slow.

For context, I am solving an ODE until a stopping criterion occurs. Then, I make some modification to the arguments of the ODE, and restart it. Schematically:

my_event = diffrax.DiscreteTerminatingEvent(lambda: state, **kwargs: my_function(state.y, *kwargs["args"]))
tcurrent = t0
y0 = my_initial_condition
args = my_initial_args
while tcurrent < t1:
    solution = diffrax.diffeqsolve(term, solver, args=args, t0=tcurrent, dt0=dt0, y0=y0,
                                   discrete_terminating_event=my_event,
                                   stepsize_controller=stepsize_controller, max_steps=None)
    tcurrent = float(solution.ts)
    y0 = solution.ys[-1]
    args = modify_args(solution.ys[-1], args)

All the functions (my_function, modify_args, the function wrapped by term) are written in JAX and JITed. The first time I run the while loop - as a cell in a jupyter notebook - it takes approx. 10 seconds. When I run it again, with identical my_initial_condition, it is significantly faster, approx. 0.3s. I assume this difference is due to the JIT compilation overhead - no problem.

However, when I re-run this with a slightly modified initial condition, e.g. y0 = my_initial_condition+1e-5 I am back to 10s runtime. This is not good, because I want to run this code block for large number of times for different values of my_initial_condition. I ran the following tests to see what might be going on:

• If the times at which the ODE solver is stopped predefined, i.e. not triggered by discrete_terminating_event, then the problem is gone, even if I am still passing a discrete_terminating_event-argument (modified so as to never trigger a stop)
• If I define my_event inside the while loop, then I always get the ~10s execution time, even if I'm re-running the cell with identical inputs. I.e.:

while tcurrent < t1:
    my_event = diffrax.DiscreteTerminatingEvent(lambda: state, **kwargs: my_function(state.y, *kwargs["args"]))
    solution = diffrax.diffeqsolve(term, solver, args=args, t0=tcurrent, dt0=dt0, y0=y0,
    ...

• When I evaluate my_function, the function inside my_event with different values of y or args, I do not trigger a JIT recompile.

This has lead my to believe that diffrax JIT-recompiles the discrete_terminating_event every time integration is stopped due to an event. Is there a way to avoid this?

Best,

Nikolas

Start with the diffrax.diffeqsolve path handling DiscreteTerminatingEvent and compare it with the predefined-stop case described in the report. Reproduce the Jupyter notebook loop using changed and unchanged y0 values, then trace JIT cache behavior around event-triggered termination; done means repeated runs avoid unnecessary recompilation while preserving event stopping behavior.