testing work

Author: lockwo

benchmarks/stateful_paths.py

@@ -163,7 +163,7 @@ def _evaluate_leaf(
 brownian_motion = diffrax.VirtualBrownianTree(t0, t1, tol=1e-3, shape=(), key=key)
 ubp = OldBrownianPath(shape=(), key=key)
 new_ubp = diffrax.UnsafeBrownianPath(shape=(), key=key)
-new_ubp_pre = diffrax.UnsafeBrownianPath(shape=(), key=key, precompute=True)
+new_ubp_pre = diffrax.UnsafeBrownianPath(shape=(), key=key, precompute=ndt + 10)
 solver = diffrax.Euler()
 terms = diffrax.MultiTerm(
     diffrax.ODETerm(drift), diffrax.ControlTerm(diffusion, brownian_motion)
@@ -177,7 +177,7 @@ def _evaluate_leaf(
 terms_new_precompute = diffrax.MultiTerm(
     diffrax.ODETerm(drift), diffrax.ControlTerm(diffusion, new_ubp_pre)
 )
-saveat = diffrax.SaveAt(ts=jnp.linspace(t0, t1, ndt))
+saveat = diffrax.SaveAt(ts=jnp.linspace(t0, t1, 1000))
 
 
 @jax.jit

diffrax/_integrate.py

@@ -387,7 +387,7 @@ def body_fun_aux(state):
 
         tprev = jnp.minimum(tprev, t1)
         tnext = _clip_to_end(tprev, tnext, t1, keep_step)
-
+        
         progress_meter_state = progress_meter.step(
             state.progress_meter_state, linear_rescale(t0, tprev, t1)
         )
@@ -862,7 +862,7 @@ class SaveAt(eqx.Module):  # noqa: F811
             t1: bool
 
 
-# @eqx.filter_jit
+@eqx.filter_jit
 @eqxi.doc_remove_args("discrete_terminating_event")
 def diffeqsolve(
     terms: PyTree[AbstractTerm],