small updates

Author: thibmonsel

diffrax/_delays.py

@@ -102,7 +102,7 @@ class Delays(eqx.Module):
     max_discontinuities: IntScalarLike = 100
     recurrent_checking: bool = False
     sub_intervals: IntScalarLike = 10
-    max_steps: IntScalarLike = 20
+    max_steps: IntScalarLike = 10
     rtol: RealScalarLike = 10e-3
     atol: RealScalarLike = 10e-3
 
@@ -136,7 +136,7 @@ def __call__(self, t, y, args):
         delays, treedef = jtu.tree_flatten(self.delays)
         if self.dense_interp is None:
             assert self.dense_info is None
-            for delay in self.delays:
+            for delay in delays:
                 delay_val = delay(t, y, args)
                 alpha_val = t - delay_val
                 y0_val = self.y0_history(alpha_val)
@@ -265,7 +265,7 @@ def fn(dense_info, args):
         atol=delays.atol,
         norm=rms_norm,
         implicit_step=implicit_step,
-        max_steps=100,
+        max_steps=delays.max_steps,
     )
 
     nonlinear_args = (

diffrax/_integrate.py

@@ -364,7 +364,7 @@ def body_fun_aux(state):
                 state.solver_state,
                 state.made_jump,
             )
-            num_dde_explicit_step = num_dde_implicit_step = 0
+            implicit_step = False
         else:
             min_delay = []
             flat_delays = jtu.tree_leaves(delays.delays)
@@ -439,64 +439,62 @@ def get_struct_dense_info(init_state):
         )
         assert jnp.result_type(keep_step) is jnp.dtype(bool)
         # Finding all of the potential discontinuity roots
-        # if delays is not None:
-        #     _part_maybe_find_discontinuity = ft.partial(
-        #         maybe_find_discontinuity,
-        #         tprev,
-        #         tnext,
-        #         dense_info,
-        #         state,
-        #         delays,
-        #         solver,
-        #         args,
-        #     )
-
-        #     tsearch = jnp.linspace(tprev, tnext, delays.sub_intervals)
-        #     batch_tprev, batch_tnext = tsearch[:-1], tsearch[1:]
-        #     vmap_maybe_find_discontinuity_wrapper = jax.vmap(
-        #         _part_maybe_find_discontinuity, (None, 0, 0)
-        #     )
-        #     if delays.recurrent_checking:
-        #         (
-        #             tnext_candidate,
-        #             batch_discont_update,
-        #         ) = vmap_maybe_find_discontinuity_wrapper(
-        #             False, batch_tprev, batch_tnext
-        #         )
-        #     else:
-        #         (
-        #             tnext_candidate,
-        #             batch_discont_update,
-        #         ) = vmap_maybe_find_discontinuity_wrapper(
-        #             keep_step, batch_tprev, batch_tnext
-        #         )
-
-        #     proxy_tnext = jnp.where(batch_discont_update, tnext_candidate, jnp.inf)
-        #     proxy_tnext = jnp.min(proxy_tnext)
-
-        #     tnext, discont_update = jax.lax.cond(
-        #         jnp.isinf(proxy_tnext),
-        #         lambda: (tnext, False),
-        #         lambda: (proxy_tnext, True),
-        #     )
-
-        #     # Count the number of steps in DDEs, just for statistical purposes
-        #     num_dde_implicit_step = state.num_dde_implicit_step + (
-        #         keep_step & implicit_step
-        #     )
-        #     num_dde_explicit_step = state.num_dde_explicit_step + (
-        #         keep_step & jnp.invert(implicit_step)
-        #     )
-
-        #     assert jnp.result_type(discont_update) is jnp.dtype(bool)
-
-        # assert jnp.result_type(keep_step) is jnp.dtype(bool)
+        discont_update = False
+        if delays is not None:
+            #     _part_maybe_find_discontinuity = ft.partial(
+            #         maybe_find_discontinuity,
+            #         tprev,
+            #         tnext,
+            #         dense_info,
+            #         state,
+            #         delays,
+            #         solver,
+            #         args,
+            #     )
+
+            #     tsearch = jnp.linspace(tprev, tnext, delays.sub_intervals)
+            #     batch_tprev, batch_tnext = tsearch[:-1], tsearch[1:]
+            #     vmap_maybe_find_discontinuity_wrapper = jax.vmap(
+            #         _part_maybe_find_discontinuity, (None, 0, 0)
+            #     )
+            #     if delays.recurrent_checking:
+            #         (
+            #             tnext_candidate,
+            #             batch_discont_update,
+            #         ) = vmap_maybe_find_discontinuity_wrapper(
+            #             False, batch_tprev, batch_tnext
+            #         )
+            #     else:
+            #         (
+            #             tnext_candidate,
+            #             batch_discont_update,
+            #         ) = vmap_maybe_find_discontinuity_wrapper(
+            #             keep_step, batch_tprev, batch_tnext
+            #         )
+
+            #     prox_tnext = jnp.where(batch_discont_update, tnext_candidate, jnp.inf)
+            #     prox_tnext = jnp.min(prox_tnext)
+
+            #     tnext, discont_update = jax.lax.cond(
+            #         jnp.isinf(prox_tnext),
+            #         lambda: (tnext, False),
+            #         lambda: (prox_tnext, True),
+            #     )
+            #     assert jnp.result_type(discont_update) is jnp.dtype(bool)
+
+            # Count the number of steps in DDEs, just for statistical purposes
+            num_dde_implicit_step = state.num_dde_implicit_step + (
+                keep_step & implicit_step
+            )
+            num_dde_explicit_step = state.num_dde_explicit_step + (
+                keep_step & jnp.invert(implicit_step)
+            )
+
+        assert jnp.result_type(keep_step) is jnp.dtype(bool)
 
         #
         # Do some book-keeping.
         #
-        discont_update = False
-        num_dde_explicit_step = num_dde_implicit_step = 0
         tprev = jnp.minimum(tprev, t1)
         tnext = _clip_to_end(tprev, tnext, t1, keep_step)
 
@@ -736,8 +734,8 @@ def _outer_cond_fn(cond_fn_i, old_event_value_i):
             event_dense_info=event_dense_info,
             event_values=event_values,
             event_mask=event_mask,
-            num_dde_explicit_step=num_dde_explicit_step,
-            num_dde_implicit_step=num_dde_implicit_step,
+            num_dde_explicit_step=num_dde_explicit_step,  # type: ignore
+            num_dde_implicit_step=num_dde_implicit_step,  # type: ignore
             discontinuities=discontinuities,  # type: ignore
             discontinuities_save_index=discontinuities_save_index,
         )