How to use DiscreteTerminatingEvent to terminate integration as soon as any ODEterm value becomes NaN?

I am using Diffrax to try to integrate magnetic field line trajectories, using a magnetic field which is defined by a spline over only in a set region of space, which will return NaN when evaluated outside of this region. I would like to set a DiscreteTerminatingEvent to stop the integration as soon as a NaN is detected in the ODEterm. How can I do this?

Note that I don't want the solver to start adaptively timestepping once NaN is detected, I simply wish to stop the integration entirely (since a NaN means the trajectory is leaving the region of validity of the magnetic field representation)

I tried doing this using

def default_terminating_event_fxn(state, **kwargs):
        terms = kwargs.get("terms", lambda a, x, b: x)
        return jnp.any(jnp.isnan(terms.vf(0, state.y, 0)))

but the solver still seems to ram into the point where NaNs occur and attempts to decrease the stepsize to the point that the max steps is reached. Is there a way for the events to access what the next value of the ODETerm will be, and act on that information?

I am unsure of exactly why this is occurring as I could not print out the values inside this event fxn I created, as even using jax debug print it only printed out tracer values and not actual numbers.

Thanks for the help!

MWE (Not a magnetic field but shows the behavior):

from diffrax import diffeqsolve, Dopri5, ODETerm, SaveAt, PIDController, DiscreteTerminatingEvent
import jax
import  jax.numpy as jnp
from jax.lax import cond
def vector_field(t,y,args):
    true_fun = lambda t,x,args: -x
    false_fun = lambda t,x,args: jnp.nan*jnp.ones_like(x)
    return cond(t<2.5, true_fun,false_fun,t,y,args)

term = ODETerm(vector_field)
solver = Dopri5()
saveat = SaveAt(ts=[0., 1., 2., 3.])
stepsize_controller = PIDController(rtol=1e-4, atol=1e-4)
def default_terminating_event_fxn(state, **kwargs):
    terms = kwargs.get("terms", lambda a, x, b: x)
    return jnp.any(jnp.isnan(terms.vf(0, state.y, 0)))
terminating_event =  DiscreteTerminatingEvent(default_terminating_event_fxn)
sol = diffeqsolve(term, solver, t0=0, t1=3, dt0=0.1, y0=1, saveat=saveat,
                  stepsize_controller=stepsize_controller, discrete_terminating_event=terminating_event)
# hits max steps and fails with XLARuntimeError, when instead expect it to complete successfully
# after exiting once NaN is reached at t=2.5


print(sol.ts)  # DeviceArray([0.   , 1.   , 2.   ,inf    ])
print(sol.ys)  # DeviceArray([1.   , 0.368, 0.135, inf])

Start by reproducing the MWE with diffeqsolve and inspect DiscreteTerminatingEvent's callback together with the solver's step and rejection path. Done means the integration terminates when the ODETerm first produces NaN, without repeated adaptive stepping or reaching the maximum step count.