From 2b6d8ff64f10898929471b4f1de21592c8a95ab4 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Wed, 21 Feb 2024 15:11:15 +0000
Subject: [PATCH 01/28] Changes to how events are handled in diffrax.

The main changes are:

    1. Added the generic Event class:
    ```
    class Event:
        event_function: PyTree[EventFn]
        root_finder: Optional[optx.AbstractRootFinder] = None
    ```
    EventFn is defined as:
    ```
    class EventFn(eqx.Module):
        cond_fn: Callable[..., Union[BoolScalarLike, RealScalarLike]]
        transition_fn: Optional[Callable[[PyTree[ArrayLike]], PyTree[ArrayLike]]] = (
            lambda x: x
        )
    ````

    2. Added root finding procedure in diffeqsolve to find exact event times that are differentiable. This is only done when root_finder is not None in the given Event class.

    3. Added event_mask to the Solution class so that, when multiple event functions are passed, the user can see which one was triggered for a given solve.

Hopefully this new event-handling is sufficiently generic to handle all kinds of events in a unified and simple manner. The main benefit is that we can now differentiate through events. So far the current implementation is enough to deal with ODEs, but I suspect more is needed for dealing with SDEs.

The new approach reduces to the old approach when passing only one EventFn with a boolean cond_fn and no transition_fn.

For now the transition_fn is not used, but it will be useful when adding non-terminating events. Similarly, we might add other attributes to EventFn to distinguish between different types of events.

No event cases in root-finding

At the end of the root-fining step (L1146 in _integrate.py), I changed:
```
return jtu.tree_map(
    _call_real,
    event.event_fn,
    final_state.event_result,
    final_state.event_compare,
    is_leaf=_is_event_fn,
)
```

to

```
results = jtu.tree_map(
    _call_real,
    event.event_fn,
    final_state.event_result,
    final_state.event_compare,
    is_leaf=_is_event_fn,
)
results_ravel, _ = jfu.ravel_pytree(results)
return jnp.where(event_happened, results_ravel, final_state.tprev - t)
```

Thus, if no event occurs the root-find will return tprev as desired. Before call_real() was constantly 0 in this case which caused in error in the root-find.

Added EventFn and Event to diffrax/__init__.py

Added tests for new event handling

I added new tests for the updated event implementation which, apart from the old ones, also checks that the right event time is found in the case where a root-find is called and that the derivatives match the theoretical derivatives.

Furthermore, I marked the following tests, that rely on the old event implementation, with @pytest.mark.skip:
 - test_event.py::test_discrete_terminate1
 - test_event.py::test_discrete_terminate2
 - test_event.py::test_event_backsolve
 - test_adjoint.py::test_implicit

In order to avoid pyright errors I had to add # pyright: ignore in a few places in the the old test referenced above.

Deleted old event implementation

I deleted the following two classes:
- diffrax._event.DiscreteTerminatingEvent
- diffrax._event.SteadyStateEvent

These were also removed from the diffrax.__init__.py

Minor changes to event hadnling

The changes are the following:

- Tweaked the event API and got rid of the EventFn class. Now there is only an Event class:

```
class Event(eqx.Module):
    cond_fn: PyTree[Callable[..., Union[BoolScalarLike, RealScalarLike]]]
    root_finder: Optional[optx.AbstractRootFinder] = None
```

- Changed the way boolean condition functions are handled in the root finding step. Now instead of calling _bool_event_gradient, we simply return result = final_state.tprev - t.

- Removed all cases where jtu.ravel_pytree was used.

- Changed "teventprev" to "tprevprev" and "event_compare" to "event_mask" in the State class.

- Updated tests.py and __init__.py to reflect the changes.

Minor changes for simplicity

I slightly changed the initialization of the event attributes in the state in _integrate.py mainly for aesthetic reasons.

Made changes according to comments on #387

No event case

Changed it so that the final value of the solve is returned in cases where no event happens instead of evaluating the interpolator.
---
 diffrax/__init__.py   |   4 +-
 diffrax/_adjoint.py   |  12 +--
 diffrax/_event.py     |  95 +-----------------
 diffrax/_integrate.py | 219 ++++++++++++++++++++++++++++++++++++++----
 diffrax/_solution.py  |   1 +
 test/test_adjoint.py  |   5 +-
 test/test_event.py    | 193 +++++++++++++++++++++++++++++++++++--
 7 files changed, 404 insertions(+), 125 deletions(-)

diff --git a/diffrax/__init__.py b/diffrax/__init__.py
index e2f094be..7426c4fc 100644
--- a/diffrax/__init__.py
+++ b/diffrax/__init__.py
@@ -22,9 +22,7 @@
     SpaceTimeTimeLevyArea as SpaceTimeTimeLevyArea,
 )
 from ._event import (
-    AbstractDiscreteTerminatingEvent as AbstractDiscreteTerminatingEvent,
-    DiscreteTerminatingEvent as DiscreteTerminatingEvent,
-    SteadyStateEvent as SteadyStateEvent,
+    Event as Event,
 )
 from ._global_interpolation import (
     AbstractGlobalInterpolation as AbstractGlobalInterpolation,
diff --git a/diffrax/_adjoint.py b/diffrax/_adjoint.py
index ff16f921..c75d0a9e 100644
--- a/diffrax/_adjoint.py
+++ b/diffrax/_adjoint.py
@@ -121,7 +121,7 @@ def loop(
         terms,
         solver,
         stepsize_controller,
-        discrete_terminating_event,
+        event,
         saveat,
         t0,
         t1,
@@ -563,7 +563,7 @@ def _loop_backsolve_bwd(
     self,
     solver,
     stepsize_controller,
-    discrete_terminating_event,
+    event,
     saveat,
     t0,
     t1,
@@ -573,7 +573,7 @@ def _loop_backsolve_bwd(
     init_state,
     progress_meter,
 ):
-    assert discrete_terminating_event is None
+    assert event is None
 
     #
     # Unpack our various arguments. Delete a lot of things just to make sure we're not
@@ -787,7 +787,7 @@ def loop(
         init_state,
         passed_solver_state,
         passed_controller_state,
-        discrete_terminating_event,
+        event,
         **kwargs,
     ):
         if jtu.tree_structure(saveat.subs, is_leaf=_is_subsaveat) != jtu.tree_structure(
@@ -829,7 +829,7 @@ def loop(
                 "`diffrax.BacksolveAdjoint` is only compatible with solvers that take "
                 "a single term."
             )
-        if discrete_terminating_event is not None:
+        if event is not None:
             raise NotImplementedError(
                 "`diffrax.BacksolveAdjoint` is not compatible with events."
             )
@@ -846,7 +846,7 @@ def loop(
             saveat=saveat,
             init_state=init_state,
             solver=solver,
-            discrete_terminating_event=discrete_terminating_event,
+            event=event,
             **kwargs,
         )
         final_state = _only_transpose_ys(final_state)
diff --git a/diffrax/_event.py b/diffrax/_event.py
index a596d8ed..7a8495b8 100644
--- a/diffrax/_event.py
+++ b/diffrax/_event.py
@@ -1,98 +1,13 @@
-import abc
 from collections.abc import Callable
-from typing import Optional
+from typing import Optional, Union
 
 import equinox as eqx
 import optimistix as optx
-from jaxtyping import Array, PyTree
+from jaxtyping import PyTree
 
 from ._custom_types import BoolScalarLike, RealScalarLike
-from ._step_size_controller import AbstractAdaptiveStepSizeController
 
 
-class AbstractDiscreteTerminatingEvent(eqx.Module):
-    """Evaluated at the end of each integration step. If true then the solve is stopped
-    at that time.
-    """
-
-    @abc.abstractmethod
-    def __call__(self, state, **kwargs) -> BoolScalarLike:
-        """**Arguments:**
-
-        - `state`: a dataclass of the evolving state of the system, including in
-            particular the solution `state.y` at time `state.tprev`.
-        - `**kwargs`: the integration options held constant throughout the solve
-            are passed as keyword arguments: `terms`, `solver`, `args`. etc.
-
-        **Returns**
-
-        A boolean. If true then the solve is terminated.
-        """
-
-
-class DiscreteTerminatingEvent(AbstractDiscreteTerminatingEvent):
-    """Terminates the solve if its condition is ever active."""
-
-    cond_fn: Callable[..., BoolScalarLike]
-
-    def __call__(self, state, **kwargs):
-        return self.cond_fn(state, **kwargs)
-
-
-DiscreteTerminatingEvent.__init__.__doc__ = """**Arguments:**
-
-- `cond_fn`: A function `(state, **kwargs) -> bool` that is evaluated on every step of
-    the differential equation solve. If it returns `True` then the solve is finished at
-    that timestep. `state` is a dataclass of the evolving state of the system,
-    including in particular the solution `state.y` at time `state.tprev`. Passed as
-    keyword arguments are the `terms`, `solver`, `args` etc. that are constant
-    throughout the solve.
-"""
-
-
-class SteadyStateEvent(AbstractDiscreteTerminatingEvent):
-    """Terminates the solve once it reaches a steady state."""
-
-    rtol: Optional[float] = None
-    atol: Optional[float] = None
-    norm: Callable[[PyTree[Array]], RealScalarLike] = optx.rms_norm
-
-    def __call__(self, state, *, terms, args, solver, stepsize_controller, **kwargs):
-        del kwargs
-        msg = (
-            "The `rtol` and `atol` tolerances for `SteadyStateEvent` default "
-            "to the `rtol` and `atol` used with an adaptive step size "
-            "controller (such as `diffrax.PIDController`). Either use an "
-            "adaptive step size controller, or specify these tolerances "
-            "manually."
-        )
-        if self.rtol is None:
-            if isinstance(stepsize_controller, AbstractAdaptiveStepSizeController):
-                _rtol = stepsize_controller.rtol
-            else:
-                raise ValueError(msg)
-        else:
-            _rtol = self.rtol
-        if self.atol is None:
-            if isinstance(stepsize_controller, AbstractAdaptiveStepSizeController):
-                _atol = stepsize_controller.atol
-            else:
-                raise ValueError(msg)
-        else:
-            _atol = self.atol
-
-        # TODO: this makes an additional function evaluation that in practice has
-        # probably already been made by the solver.
-        vf = solver.func(terms, state.tprev, state.y, args)
-        return self.norm(vf) < _atol + _rtol * self.norm(state.y)
-
-
-SteadyStateEvent.__init__.__doc__ = """**Arguments:**
-
-- `rtol`: The relative tolerance for determining convergence. Defaults to the
-    same `rtol` as passed to an adaptive step controller if one is used.
-- `atol`: The absolute tolerance for determining convergence. Defaults to the
-    same `atol` as passed to an adaptive step controller if one is used.
-- `norm`: A function `PyTree -> Scalar`, which is called to determine whether
-    the vector field is close to zero.
-"""
+class Event(eqx.Module):
+    cond_fn: PyTree[Callable[..., Union[BoolScalarLike, RealScalarLike]]]
+    root_finder: Optional[optx.AbstractRootFinder] = None
diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index 6cb25140..e4824bef 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -9,6 +9,7 @@
     Optional,
     Tuple,
     TYPE_CHECKING,
+    Union,
 )
 
 import equinox as eqx
@@ -18,17 +19,19 @@
 import jax.numpy as jnp
 import jax.tree_util as jtu
 import lineax.internal as lxi
+import optimistix as optx
 from jaxtyping import Array, ArrayLike, Float, Inexact, PyTree, Real
 
 from ._adjoint import AbstractAdjoint, RecursiveCheckpointAdjoint
 from ._custom_types import (
     BoolScalarLike,
     BufferDenseInfos,
+    DenseInfo,
     FloatScalarLike,
     IntScalarLike,
     RealScalarLike,
 )
-from ._event import AbstractDiscreteTerminatingEvent
+from ._event import Event
 from ._global_interpolation import DenseInterpolation
 from ._heuristics import is_sde, is_unsafe_sde
 from ._misc import linear_rescale, static_select
@@ -84,12 +87,20 @@ class State(eqx.Module):
     dense_infos: Optional[BufferDenseInfos]
     dense_save_index: Optional[IntScalarLike]
     progress_meter_state: PyTree[Array]
+    event_result: Optional[PyTree[Union[BoolScalarLike, RealScalarLike]]]
+    event_mask: Optional[PyTree[BoolScalarLike]]
+    dense_info_for_event: Optional[DenseInfo]
+    tprevprev: Optional[FloatScalarLike]
 
 
 def _is_none(x: Any) -> bool:
     return x is None
 
 
+def _is_bool(x: Any) -> bool:
+    return isinstance(x, BoolScalarLike)
+
+
 def _term_compatible(
     y: PyTree[ArrayLike],
     args: PyTree[Any],
@@ -236,6 +247,36 @@ def _maybe_static(static_x: Optional[ArrayLike], x: ArrayLike) -> ArrayLike:
         return x
 
 
+def _compare_events(
+    old_event_result: Union[BoolScalarLike, RealScalarLike],
+    new_event_result: Union[BoolScalarLike, RealScalarLike],
+) -> BoolScalarLike:
+    with jax.numpy_dtype_promotion("standard"):
+        old_dtype = jnp.result_type(jnp.array(old_event_result), jnp.float32)
+        new_dtype = jnp.result_type(jnp.array(new_event_result), jnp.float32)
+        old_sign = jnp.sign(jnp.array(old_event_result, dtype=old_dtype))
+        new_sign = jnp.sign(jnp.array(new_event_result, dtype=new_dtype))
+    return old_sign != new_sign
+
+
+def _event_happened(event_mask: PyTree[BoolScalarLike]) -> BoolScalarLike:
+    return jnp.any(jnp.array(jtu.tree_leaves(event_mask)))
+
+
+@jax.custom_jvp
+def _bool_event_gradient(t: RealScalarLike, result: BoolScalarLike) -> FloatScalarLike:
+    return jnp.where(result, 0.0, 1.0)
+
+
+@_bool_event_gradient.defjvp
+def _bool_event_gradient_jvp(primals, tangents):
+    t, result = primals
+    tangent_t, _ = tangents
+    out = _bool_event_gradient(t, result)
+    tangent_out = 1.0
+    return out, tangent_out
+
+
 _PRINT_STATIC = False  # used in tests
 
 
@@ -243,7 +284,7 @@ def loop(
     *,
     solver,
     stepsize_controller,
-    discrete_terminating_event,
+    event,
     saveat,
     t0,
     t1,
@@ -318,6 +359,14 @@ def body_fun_aux(state):
         # everything breaks.) See #143.
         y_error = jtu.tree_map(lambda x: jnp.where(jnp.isnan(x), jnp.inf, x), y_error)
 
+        # Save info for event handling
+        if event is None:
+            tprevprev = None
+            dense_info_for_event = None
+        else:
+            tprevprev = state.tprev
+            dense_info_for_event = dense_info
+
         error_order = solver.error_order(terms)
         (
             keep_step,
@@ -477,27 +526,56 @@ def save_steps(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState:
             dense_infos=dense_infos,
             dense_save_index=dense_save_index,
             progress_meter_state=progress_meter_state,
+            dense_info_for_event=dense_info_for_event,
+            tprevprev=tprevprev,
+            event_result=None,
+            event_mask=None,
         )
 
-        if discrete_terminating_event is not None:
-            discrete_terminating_event_occurred = discrete_terminating_event(
-                new_state,
-                solver=solver,
-                stepsize_controller=stepsize_controller,
-                saveat=saveat,
-                t0=t0,
-                t1=t1,
-                dt0=dt0,
-                max_steps=max_steps,
-                terms=terms,
-                args=args,
-            )
+        if event is not None:
+
+            def _call_event(_cond_fn):
+                return _cond_fn(
+                    new_state,
+                    y=y,
+                    solver=solver,
+                    stepsize_controller=stepsize_controller,
+                    saveat=saveat,
+                    t0=t0,
+                    t1=t1,
+                    dt0=dt0,
+                    max_steps=max_steps,
+                    terms=terms,
+                    args=args,
+                )
+
+            event_result = jtu.tree_map(_call_event, event.cond_fn, is_leaf=callable)
+            event_mask = jtu.tree_map(_compare_events, state.event_result, event_result)
             result = RESULTS.where(
-                discrete_terminating_event_occurred,
+                _event_happened(event_mask),
                 RESULTS.discrete_terminating_event_occurred,
                 result,
             )
+            # If multiple events are triggered in one step we take the first one
+            event_hist = []
+
+            def counter_update(x):
+                out = jnp.where(jnp.any(jnp.array(event_hist)), False, x)
+                event_hist.append(x)
+                return out
+
+            event_mask = jtu.tree_map(counter_update, event_mask, is_leaf=_is_bool)
+
             new_state = eqx.tree_at(lambda s: s.result, new_state, result)
+            new_state = eqx.tree_at(
+                lambda s: s.event_result, new_state, event_result, is_leaf=_is_none
+            )
+            new_state = eqx.tree_at(
+                lambda s: s.event_mask,
+                new_state,
+                event_mask,
+                is_leaf=_is_none,
+            )
 
         return (
             new_state,
@@ -578,7 +656,7 @@ def diffeqsolve(
     saveat: SaveAt = SaveAt(t1=True),
     stepsize_controller: AbstractStepSizeController = ConstantStepSize(),
     adjoint: AbstractAdjoint = RecursiveCheckpointAdjoint(),
-    discrete_terminating_event: Optional[AbstractDiscreteTerminatingEvent] = None,
+    event: Optional[Event] = None,
     max_steps: Optional[int] = 4096,
     throw: bool = True,
     progress_meter: AbstractProgressMeter = NoProgressMeter(),
@@ -977,6 +1055,22 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
         dense_infos = None
         dense_save_index = None
 
+    if event is None:
+        tprevprev = None
+        dense_info_for_event = None
+        event_mask = None
+    else:
+        tprevprev = tprev
+        _, _, traced_dense_info_event, _, _ = eqx.filter_eval_shape(
+            solver.step, terms, tprev, tnext, y0, args, solver_state, made_jump
+        )
+        dense_info_for_event = jtu.tree_map(
+            lambda x: jnp.empty(x.shape, x.dtype), traced_dense_info_event
+        )
+        event_mask = jtu.tree_map(
+            lambda x: jnp.bool_(False), event.cond_fn, is_leaf=callable
+        )
+
     # Progress meter
     progress_meter_state = progress_meter.init()
 
@@ -997,8 +1091,35 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
         dense_infos=dense_infos,
         dense_save_index=dense_save_index,
         progress_meter_state=progress_meter_state,
+        tprevprev=tprevprev,
+        dense_info_for_event=dense_info_for_event,
+        event_result=None,
+        event_mask=event_mask,
     )
 
+    # Since cond_fn might depend on the state, we need to initialise event_result here.
+    if event is not None:
+
+        def _call_event(_cond_fn):
+            return _cond_fn(
+                init_state,
+                y=y0,
+                solver=solver,
+                stepsize_controller=stepsize_controller,
+                saveat=saveat,
+                t0=t0,
+                t1=t1,
+                dt0=dt0,
+                max_steps=max_steps,
+                terms=terms,
+                args=args,
+            )
+
+        event_result = jtu.tree_map(_call_event, event.cond_fn, is_leaf=callable)
+        init_state = eqx.tree_at(
+            lambda s: s.event_result, init_state, event_result, is_leaf=_is_none
+        )
+
     #
     # Main loop
     #
@@ -1008,7 +1129,7 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
         terms=terms,
         solver=solver,
         stepsize_controller=stepsize_controller,
-        discrete_terminating_event=discrete_terminating_event,
+        event=event,
         saveat=saveat,
         t0=t0,
         t1=t1,
@@ -1031,6 +1152,67 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
         lambda s: s.ts * direction, final_state.save_state, is_leaf=is_save_state
     )
     ys = jtu.tree_map(lambda s: s.ys, final_state.save_state, is_leaf=is_save_state)
+
+    # Do root find for exact event times
+    if event is not None:
+        event_mask = final_state.event_mask
+        event_happened = _event_happened(event_mask)
+        interpolator = solver.interpolation_cls(
+            t0=final_state.tprevprev,
+            t1=final_state.tprev,
+            **final_state.dense_info_for_event,
+        )
+        if event.root_finder is None:
+            tevent = final_state.tprev
+        else:
+
+            def _to_root_find(t, args):
+                def _call_real(cond_fn_i, event_mask_i):
+                    y = interpolator.evaluate(t)
+                    result = cond_fn_i(
+                        final_state,
+                        y=y,
+                        solver=solver,
+                        stepsize_controller=stepsize_controller,
+                        saveat=saveat,
+                        t0=t0,
+                        t1=t1,
+                        dt0=dt0,
+                        max_steps=max_steps,
+                        terms=terms,
+                        args=args,
+                    )
+
+                    if jnp.result_type(result) == jnp.bool_:
+                        result = final_state.tprev - t
+                    return jnp.where(event_mask_i, result, 0.0)
+
+                results = jtu.tree_map(
+                    _call_real,
+                    event.cond_fn,
+                    event_mask,
+                    is_leaf=callable,
+                )
+                # If no events are triggered simply push tevent towards tprev
+                results = jtu.tree_map(
+                    lambda x: jnp.where(event_happened, x, final_state.tprev - t),
+                    results,
+                )
+                return results
+
+            options = {"lower": final_state.tprevprev, "upper": final_state.tprev}
+            roots = optx.root_find(
+                _to_root_find, event.root_finder, y0=final_state.tprev, options=options
+            )
+            tevent = roots.value
+
+        # We might need to change this in order to get more accurate derivatives
+        yevent = jnp.where(event_happened, interpolator.evaluate(tevent), ys[-1])
+        ys = jtu.tree_map(lambda _y, _yevent: _y.at[-1].set(_yevent), ys, yevent)
+        ts = ts.at[-1].set(tevent)
+    else:
+        event_mask = None
+
     # It's important that we don't do any further postprocessing on `ys` here, as
     # it is the `final_state` value that is used when backpropagating via
     # optimise-then-discretise.
@@ -1082,6 +1264,7 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
         solver_state=solver_state,
         controller_state=controller_state,
         made_jump=made_jump,
+        event_mask=event_mask,
     )
 
     if throw:
diff --git a/diffrax/_solution.py b/diffrax/_solution.py
index e26a303b..35fcd51e 100644
--- a/diffrax/_solution.py
+++ b/diffrax/_solution.py
@@ -99,6 +99,7 @@ class Solution(AbstractPath):
     solver_state: Optional[PyTree]
     controller_state: Optional[PyTree]
     made_jump: Optional[BoolScalarLike]
+    event_mask: Optional[PyTree[BoolScalarLike]]
 
     def evaluate(
         self, t0: RealScalarLike, t1: Optional[RealScalarLike] = None, left: bool = True
diff --git a/test/test_adjoint.py b/test/test_adjoint.py
index f61a3f5e..f0a15e71 100644
--- a/test/test_adjoint.py
+++ b/test/test_adjoint.py
@@ -258,6 +258,7 @@ def run(model):
     run(mlp)
 
 
+@pytest.mark.skip(reason="SteadyStateEvent discontinued")
 def test_implicit():
     class ExponentialDecayToSteadyState(eqx.Module):
         steady_state: Array
@@ -275,7 +276,7 @@ def loss(model, target_steady_state):
         y0 = 1.0
         max_steps = None
         controller = diffrax.PIDController(rtol=1e-3, atol=1e-6)
-        event = diffrax.SteadyStateEvent()
+        event = diffrax.SteadyStateEvent()  # pyright: ignore
         adjoint = diffrax.ImplicitAdjoint()
         sol = diffrax.diffeqsolve(
             term,
@@ -286,7 +287,7 @@ def loss(model, target_steady_state):
             y0,
             max_steps=max_steps,
             stepsize_controller=controller,
-            discrete_terminating_event=event,
+            discrete_terminating_event=event,  # pyright: ignore
             adjoint=adjoint,
         )
         (y1,) = cast(Array, sol.ys)
diff --git a/test/test_event.py b/test/test_event.py
index 135ba4de..64a54822 100644
--- a/test/test_event.py
+++ b/test/test_event.py
@@ -3,10 +3,12 @@
 import diffrax
 import jax
 import jax.numpy as jnp
+import optimistix as optx
 import pytest
 from jaxtyping import Array
 
 
+@pytest.mark.skip(reason="Old event implementation")
 def test_discrete_terminate1():
     term = diffrax.ODETerm(lambda t, y, args: y)
     solver = diffrax.Tsit5()
@@ -19,13 +21,20 @@ def event_fn(state, **kwargs):
         assert isinstance(state.y, jax.Array)
         return state.tprev > 10
 
-    event = diffrax.DiscreteTerminatingEvent(event_fn)
+    event = diffrax.DiscreteTerminatingEvent(event_fn)  # pyright: ignore
     sol = diffrax.diffeqsolve(
-        term, solver, t0, t1, dt0, y0, discrete_terminating_event=event
+        term,
+        solver,
+        t0,
+        t1,
+        dt0,
+        y0,
+        discrete_terminating_event=event,  # pyright: ignore
     )
     assert jnp.all(cast(Array, sol.ys) > 10)
 
 
+@pytest.mark.skip(reason="Old event implementation")
 def test_discrete_terminate2():
     term = diffrax.ODETerm(lambda t, y, args: y)
     solver = diffrax.Tsit5()
@@ -38,13 +47,20 @@ def event_fn(state, **kwargs):
         assert isinstance(state.y, jax.Array)
         return state.tprev > 10
 
-    event = diffrax.DiscreteTerminatingEvent(event_fn)
+    event = diffrax.DiscreteTerminatingEvent(event_fn)  # pyright: ignore
     sol = diffrax.diffeqsolve(
-        term, solver, t0, t1, dt0, y0, discrete_terminating_event=event
+        term,
+        solver,
+        t0,
+        t1,
+        dt0,
+        y0,
+        discrete_terminating_event=event,  # pyright: ignore
     )
     assert jnp.all(cast(Array, sol.ts) > 10)
 
 
+@pytest.mark.skip(reason="Old event implementation")
 def test_event_backsolve():
     term = diffrax.ODETerm(lambda t, y, args: y)
     solver = diffrax.Tsit5()
@@ -57,7 +73,7 @@ def event_fn(state, **kwargs):
         assert isinstance(state.y, jax.Array)
         return state.tprev > 10
 
-    event = diffrax.DiscreteTerminatingEvent(event_fn)
+    event = diffrax.DiscreteTerminatingEvent(event_fn)  # pyright: ignore
 
     @jax.jit
     @jax.grad
@@ -69,7 +85,7 @@ def run(y0):
             t1,
             dt0,
             y0,
-            discrete_terminating_event=event,
+            discrete_terminating_event=event,  # pyright: ignore
             adjoint=diffrax.BacksolveAdjoint(),
         )
         return jnp.sum(cast(Array, sol.ys))
@@ -80,3 +96,168 @@ def run(y0):
 
 
 # diffrax.SteadyStateEvent tested as part of test_adjoint.py::test_implicit
+
+
+def test_continuous_terminate1():
+    term = diffrax.ODETerm(lambda t, y, args: y)
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = jnp.inf
+    dt0 = 1
+    y0 = 1.0
+
+    def cond_fn(state, **kwargs):
+        assert isinstance(state.y, jax.Array)
+        return state.tprev > 10
+
+    event = diffrax.Event(cond_fn=cond_fn)
+    sol = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
+    assert jnp.all(cast(Array, sol.ys) > 10)
+
+
+def test_continuous_terminate2():
+    term = diffrax.ODETerm(lambda t, y, args: y)
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = jnp.inf
+    dt0 = 1
+    y0 = 1.0
+
+    def cond_fn(state, **kwargs):
+        assert isinstance(state.y, jax.Array)
+        return state.tprev - 10.0
+
+    event = diffrax.Event(cond_fn=cond_fn)
+    sol = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
+    assert jnp.all(cast(Array, sol.ts) >= 10)
+
+
+def test_continuous_event_time():
+    term = diffrax.ODETerm(lambda t, y, args: 1.0)
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = jnp.inf
+    dt0 = 1.0
+    y0 = -10.0
+
+    def cond_fn(state, y, **kwargs):
+        assert isinstance(state.y, jax.Array)
+        return y
+
+    root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
+    event = diffrax.Event(cond_fn, root_finder)
+    sol = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
+    assert jnp.all(jnp.isclose(cast(Array, sol.ts), 10.0, 1e-5))
+
+
+def test_continuous_event_value():
+    term = diffrax.ODETerm(lambda t, y, args: 1.0)
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = jnp.inf
+    dt0 = 1.0
+    y0 = -10.0
+
+    def cond_fn(state, y, **kwargs):
+        assert isinstance(state.y, jax.Array)
+        return y
+
+    root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
+    event = diffrax.Event(cond_fn, root_finder)
+    sol = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
+    assert jnp.all(jnp.isclose(cast(Array, sol.ys), 0.0, 1e-5))
+
+
+def test_continuous_no_event():
+    term = diffrax.ODETerm(lambda t, y, args: 1.0)
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = 5.0
+    dt0 = 1.0
+    y0 = -10.0
+
+    def cond_fn(state, y, **kwargs):
+        assert isinstance(state.y, jax.Array)
+        return y
+
+    root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
+    event = diffrax.Event(cond_fn, root_finder)
+    sol = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
+    assert not cast(Array, sol.event_mask)
+    assert jnp.all(jnp.isclose(cast(Array, sol.ts), 5.0, 1e-5))
+
+
+def test_continuous_two_events():
+    term = diffrax.ODETerm(lambda t, y, args: 1.0)
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = jnp.inf
+    dt0 = 1.0
+    y0 = -10.0
+
+    def cond_fn_1(state, y, **kwargs):
+        assert isinstance(state.y, jax.Array)
+        return y
+
+    def cond_fn_2(state, y, **kwargs):
+        assert isinstance(state.y, jax.Array)
+        return y + 5.0
+
+    root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
+    event = diffrax.Event([cond_fn_1, cond_fn_2], root_finder)
+    sol = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
+    assert cast(Array, sol.event_mask)[1]
+    assert not cast(Array, sol.event_mask)[0]
+    assert jnp.all(jnp.isclose(cast(Array, sol.ts), 5.0, 1e-5))
+
+
+def test_continuous_event_time_grad():
+    def vector_field(t, y, args):
+        x, v = y
+        d_out = v, -8.0
+        return jnp.array(d_out)
+
+    def cond_fn(state, y, **kwargs):
+        x, v = y
+        return x
+
+    term = diffrax.ODETerm(vector_field)
+    solver = diffrax.Tsit5()
+    root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
+    event = diffrax.Event(cond_fn, root_finder)
+    t0 = 0
+    t1 = jnp.inf
+    dt0 = 0.01
+
+    @jax.jit
+    @jax.grad
+    def first_bounce_time(x0):
+        y0 = jnp.array([x0, 0.0])
+        sol = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
+        assert sol.ts is not None
+        return sol.ts[-1]
+
+    def first_bounce_time_grad(x0):
+        y0 = jnp.array([x0, 0.0])
+        sol0 = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
+        assert sol0.ts is not None
+        tevent = jax.lax.stop_gradient(sol0.ts[-1])
+
+        def phi(_x):
+            _y0 = jnp.array([_x, 0.0])
+            _sol = diffrax.diffeqsolve(term, solver, t0, tevent, dt0, _y0)
+            assert _sol.ys is not None
+            return _sol.ys[-1, :]
+
+        def event_fn(_y):
+            return cond_fn(None, _y)
+
+        _, num = jax.vjp(phi, x0)
+        (num,) = num(jax.grad(event_fn)(y0))
+        _, dem = jax.jvp(event_fn, (y0,), (vector_field(tevent, phi(x0), None),))
+        return -num / dem
+
+    x0_test = jnp.array([1.0, 3.0, 10.0])
+    x0_autograd = jax.vmap(first_bounce_time)(x0_test)
+    x0_truegrad = jax.vmap(first_bounce_time_grad)(x0_test)
+    assert jnp.all(jnp.isclose(x0_autograd, x0_truegrad, 1e-5))

From 3e367f272672453bdd0d8427d6435e962dcc7325 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Tue, 7 May 2024 11:41:22 +0200
Subject: [PATCH 02/28] Test now fails when no root finder is provided

---
 test/test_event.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/test/test_event.py b/test/test_event.py
index 64a54822..d9334e08 100644
--- a/test/test_event.py
+++ b/test/test_event.py
@@ -133,21 +133,21 @@ def cond_fn(state, **kwargs):
 
 
 def test_continuous_event_time():
-    term = diffrax.ODETerm(lambda t, y, args: 1.0)
+    term = diffrax.ODETerm(lambda t, y, args: y)
     solver = diffrax.Tsit5()
     t0 = 0
     t1 = jnp.inf
     dt0 = 1.0
-    y0 = -10.0
+    y0 = 1.0
 
     def cond_fn(state, y, **kwargs):
         assert isinstance(state.y, jax.Array)
-        return y
+        return y - jnp.exp(1.0)
 
     root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
     event = diffrax.Event(cond_fn, root_finder)
     sol = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
-    assert jnp.all(jnp.isclose(cast(Array, sol.ts), 10.0, 1e-5))
+    assert jnp.all(jnp.isclose(cast(Array, sol.ts), 1.0, 1e-4))
 
 
 def test_continuous_event_value():

From 3f67a27c67dcfb8bc98c9a0b92789e2ee2e54553 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Wed, 15 May 2024 09:24:27 +0200
Subject: [PATCH 03/28] Saving events with `SubSaveAt`s

Previously, updating the last element of ys and ts did not handle the case where multiple `SubSaveAt`s were used. This is now fixed by adding a `jtu.tree_map` in the appropriate place.
---
 diffrax/_integrate.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index e4824bef..7c2928b8 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -1207,9 +1207,12 @@ def _call_real(cond_fn_i, event_mask_i):
             tevent = roots.value
 
         # We might need to change this in order to get more accurate derivatives
-        yevent = jnp.where(event_happened, interpolator.evaluate(tevent), ys[-1])
+        yevent = jtu.tree_map(
+            lambda _y: jnp.where(event_happened, interpolator.evaluate(tevent), _y[-1]),
+            ys,
+        )
         ys = jtu.tree_map(lambda _y, _yevent: _y.at[-1].set(_yevent), ys, yevent)
-        ts = ts.at[-1].set(tevent)
+        ts = jtu.tree_map(lambda _t: _t.at[-1].set(tevent), ts)
     else:
         event_mask = None
 

From 76762c1c8dbc63513248aa693060fb392e5f1dc0 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Wed, 15 May 2024 19:55:06 +0200
Subject: [PATCH 04/28] Accounting for `SubSaveAt.fn` returning a PyTree

---
 diffrax/_integrate.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index 7c2928b8..f9a381c8 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -1207,11 +1207,14 @@ def _call_real(cond_fn_i, event_mask_i):
             tevent = roots.value
 
         # We might need to change this in order to get more accurate derivatives
-        yevent = jtu.tree_map(
-            lambda _y: jnp.where(event_happened, interpolator.evaluate(tevent), _y[-1]),
-            ys,
-        )
-        ys = jtu.tree_map(lambda _y, _yevent: _y.at[-1].set(_yevent), ys, yevent)
+        yevent = interpolator.evaluate(tevent)
+
+        def _save_yevent(subsaveat: SubSaveAt, y):
+            return jtu.tree_map(
+                lambda s, _y: _y.at[-1].set(s), subsaveat.fn(tevent, yevent, args), y
+            )
+
+        ys = jtu.tree_map(_save_yevent, saveat.subs, ys, is_leaf=_is_subsaveat)
         ts = jtu.tree_map(lambda _t: _t.at[-1].set(tevent), ts)
     else:
         event_mask = None

From 76cd083613295721533124ac6b8883f7b24f9011 Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Sun, 19 May 2024 20:44:27 +0200
Subject: [PATCH 05/28] Adjustments to #387 (events):

- Semantic change: boolean events now trigger when they become truthy (before they occurred when they swap being falsy<->truthy). Note that this required twiddling around a few things as previously it was impossible for an event to trigger on the first step; now they can.
- Semantic change: event functions now have the signature `(t, y, args *, terms, solver, **etc)` for consistency with vector fields and with `SaveAt(fn=...)`.
- Feature: now backward-compatible with the old discrete terminating events.
- Feature: added `diffrax.steady_state_event`.
- Bugfix: the final `t` and `y` from an event are now saved in the correct index of `ts` and `ys`, rather than just always being saved at index `-1`.
- Bugfix: at one point `args` referred to the `args` coming from a root find rather than the overall `diffeqsolve`.
- Bugfix: the current `state.tprev` was used instead of the previous state's `tnext`. (These are usually but not always the same -- in particular when around jumps.)
- Bugfix: added some checks when the condition function of an event does not return a bool/float scalar.
- Performance: includes a fastpath for skipping the rootfind if no events are triggered.
- Performance: now avoiding tracing for the shape of `dense_info` twice when using adaptive step size controllers alongside events.
- Performance: avoided quadratic loop for figuring out what was the first event to trigger.
- Chore: added support for the possibility of the final root find (for the time of the event) failing.
- Chore: removed some dead code (`_bool_event_gradient`).
- Chore: removed references in the docs to the old `discrete_terminating_event`.

In addition, some drive-bys:

- Fixed warnings about pending deprecations `jnp.clip(..., a_min=..., a_max=...)`.
- Had `aux_stats` (in `_integrate.py`) forward to the overall output statistics. In practice this is empty but it's worth doing for the future.
---
 diffrax/__init__.py   |   6 +
 diffrax/_event.py     | 107 ++++++++-
 diffrax/_integrate.py | 534 ++++++++++++++++++++++++------------------
 diffrax/_solution.py  |  27 ++-
 docs/api/events.md    |  20 +-
 test/test_adjoint.py  |   6 +-
 test/test_event.py    | 120 +++++-----
 7 files changed, 512 insertions(+), 308 deletions(-)

diff --git a/diffrax/__init__.py b/diffrax/__init__.py
index 7426c4fc..67b4ca50 100644
--- a/diffrax/__init__.py
+++ b/diffrax/__init__.py
@@ -22,7 +22,13 @@
     SpaceTimeTimeLevyArea as SpaceTimeTimeLevyArea,
 )
 from ._event import (
+    # Deliberately not provided with `X as X` as these are now deprecated, so we'd like
+    # static type checkers to warn about using them.
+    AbstractDiscreteTerminatingEvent,  # noqa: F401
+    DiscreteTerminatingEvent,  # noqa: F401
     Event as Event,
+    steady_state_event as steady_state_event,
+    SteadyStateEvent,  # noqa: F401
 )
 from ._global_interpolation import (
     AbstractGlobalInterpolation as AbstractGlobalInterpolation,
diff --git a/diffrax/_event.py b/diffrax/_event.py
index 7a8495b8..6f3cbc58 100644
--- a/diffrax/_event.py
+++ b/diffrax/_event.py
@@ -1,13 +1,116 @@
+import abc
 from collections.abc import Callable
 from typing import Optional, Union
 
 import equinox as eqx
 import optimistix as optx
-from jaxtyping import PyTree
+from jaxtyping import Array, PyTree
 
-from ._custom_types import BoolScalarLike, RealScalarLike
+from ._custom_types import BoolScalarLike, FloatScalarLike, RealScalarLike
+from ._step_size_controller import AbstractAdaptiveStepSizeController
 
 
 class Event(eqx.Module):
     cond_fn: PyTree[Callable[..., Union[BoolScalarLike, RealScalarLike]]]
     root_finder: Optional[optx.AbstractRootFinder] = None
+
+
+def steady_state_event(
+    rtol: Optional[float] = None,
+    atol: Optional[float] = None,
+    norm: Optional[Callable[[PyTree[Array]], RealScalarLike]] = None,
+):
+    """Create a [`diffrax.Event`][] that terminates the solve once a steady state is
+    achieved.
+
+    **Arguments:**
+
+    - `rtol`, `atol`, `norm`: the solve will terminate once
+        `norm(f) < atol + rtol * norm(y)`, where `f` is the result of evaluating the
+        vector field. Will default to the values used in the `stepsize_controller` if
+        they are not specified here.
+
+    **Returns:**
+
+    A [`diffrax.Event`][] object, that can be passed to
+    `diffrax.diffeqsolve(..., event=...)`.
+    """
+
+    def _cond_fn(t, y, args, *, terms, solver, stepsize_controller, **kwargs):
+        del kwargs
+        msg = (
+            "The `rtol`, `atol`, and `norm` for `steady_state_event` default to the "
+            "values used with an adaptive step size controller (such as "
+            "`diffrax.PIDController`). Either use an adaptive step size controller, or "
+            "specify these tolerances manually."
+        )
+        if rtol is None:
+            if isinstance(stepsize_controller, AbstractAdaptiveStepSizeController):
+                _rtol = stepsize_controller.rtol
+            else:
+                raise ValueError(msg)
+        else:
+            _rtol = rtol
+        if atol is None:
+            if isinstance(stepsize_controller, AbstractAdaptiveStepSizeController):
+                _atol = stepsize_controller.atol
+            else:
+                raise ValueError(msg)
+        else:
+            _atol = atol
+        if norm is None:
+            if isinstance(stepsize_controller, AbstractAdaptiveStepSizeController):
+                _norm = stepsize_controller.norm
+            else:
+                raise ValueError(msg)
+        else:
+            _norm = norm
+
+        # TODO: this makes an additional function evaluation that in practice has
+        # probably already been made by the solver.
+        vf = solver.func(terms, t, y, args)
+        return _norm(vf) < _atol + _rtol * _norm(y)
+
+    return Event(cond_fn=_cond_fn)
+
+
+#
+# Backward compatibility: continue to support `AbstractDiscreteTerminatingEvent`.
+# TODO: eventually remove everything below this line.
+#
+
+
+class AbstractDiscreteTerminatingEvent(eqx.Module):
+    @abc.abstractmethod
+    def __call__(self, state, **kwargs) -> BoolScalarLike:
+        pass
+
+
+class DiscreteTerminatingEvent(AbstractDiscreteTerminatingEvent):
+    cond_fn: Callable[..., BoolScalarLike]
+
+    def __call__(self, state, **kwargs):
+        return self.cond_fn(state, **kwargs)
+
+
+class SteadyStateEvent(AbstractDiscreteTerminatingEvent):
+    rtol: Optional[float] = None
+    atol: Optional[float] = None
+    norm: Callable[[PyTree[Array]], RealScalarLike] = optx.rms_norm
+
+    def __call__(self, state, *, args, **kwargs):
+        return steady_state_event(self.rtol, self.atol, self.norm).cond_fn(
+            state.tprev, state.y, args, **kwargs
+        )
+
+
+class _StateCompat(eqx.Module):
+    tprev: FloatScalarLike
+    y: PyTree[Array]
+
+
+class DiscreteTerminatingEventToCondFn(eqx.Module):
+    event: AbstractDiscreteTerminatingEvent
+
+    def __call__(self, t, y, args, **kwargs):
+        return self.event(_StateCompat(tprev=t, y=y), args=args, **kwargs)
diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index f9a381c8..0b58e7c7 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -16,6 +16,7 @@
 import equinox.internal as eqxi
 import jax
 import jax.core
+import jax.lax as lax
 import jax.numpy as jnp
 import jax.tree_util as jtu
 import lineax.internal as lxi
@@ -31,7 +32,11 @@
     IntScalarLike,
     RealScalarLike,
 )
-from ._event import Event
+from ._event import (
+    AbstractDiscreteTerminatingEvent,
+    DiscreteTerminatingEventToCondFn,
+    Event,
+)
 from ._global_interpolation import DenseInterpolation
 from ._heuristics import is_sde, is_unsafe_sde
 from ._misc import linear_rescale, static_select
@@ -70,37 +75,47 @@ class SaveState(eqx.Module):
 
 
 class State(eqx.Module):
+    #
     # Evolving state during the solve
+    #
     y: PyTree[Array]
     tprev: FloatScalarLike
     tnext: FloatScalarLike
     made_jump: BoolScalarLike
     solver_state: PyTree[ArrayLike]
     controller_state: PyTree[ArrayLike]
+    progress_meter_state: PyTree[Array]
     result: RESULTS
+    #
+    # Reported output statistics
+    #
     num_steps: IntScalarLike
     num_accepted_steps: IntScalarLike
     num_rejected_steps: IntScalarLike
+    #
     # Output that is .at[].set() updated during the solve (and their indices)
+    #
     save_state: PyTree[SaveState]
     dense_ts: Optional[eqxi.MaybeBuffer[Float[Array, " times_plus_1"]]]
     dense_infos: Optional[BufferDenseInfos]
     dense_save_index: Optional[IntScalarLike]
-    progress_meter_state: PyTree[Array]
-    event_result: Optional[PyTree[Union[BoolScalarLike, RealScalarLike]]]
+    #
+    # Information about the most recent step, used for events.
+    #
+    # Not recorded anywhere else: this is the previous state's `tprev`.
+    event_tprev: Optional[FloatScalarLike]
+    # This is the previous state's `tnext`. This is not necessarily the same as our
+    # `tprev`, as the two can differ a little bit when crossing jumps.
+    event_tnext: Optional[FloatScalarLike]
+    event_dense_info: Optional[DenseInfo]
+    event_values: Optional[PyTree[Union[BoolScalarLike, RealScalarLike]]]
     event_mask: Optional[PyTree[BoolScalarLike]]
-    dense_info_for_event: Optional[DenseInfo]
-    tprevprev: Optional[FloatScalarLike]
 
 
 def _is_none(x: Any) -> bool:
     return x is None
 
 
-def _is_bool(x: Any) -> bool:
-    return isinstance(x, BoolScalarLike)
-
-
 def _term_compatible(
     y: PyTree[ArrayLike],
     args: PyTree[Any],
@@ -247,36 +262,6 @@ def _maybe_static(static_x: Optional[ArrayLike], x: ArrayLike) -> ArrayLike:
         return x
 
 
-def _compare_events(
-    old_event_result: Union[BoolScalarLike, RealScalarLike],
-    new_event_result: Union[BoolScalarLike, RealScalarLike],
-) -> BoolScalarLike:
-    with jax.numpy_dtype_promotion("standard"):
-        old_dtype = jnp.result_type(jnp.array(old_event_result), jnp.float32)
-        new_dtype = jnp.result_type(jnp.array(new_event_result), jnp.float32)
-        old_sign = jnp.sign(jnp.array(old_event_result, dtype=old_dtype))
-        new_sign = jnp.sign(jnp.array(new_event_result, dtype=new_dtype))
-    return old_sign != new_sign
-
-
-def _event_happened(event_mask: PyTree[BoolScalarLike]) -> BoolScalarLike:
-    return jnp.any(jnp.array(jtu.tree_leaves(event_mask)))
-
-
-@jax.custom_jvp
-def _bool_event_gradient(t: RealScalarLike, result: BoolScalarLike) -> FloatScalarLike:
-    return jnp.where(result, 0.0, 1.0)
-
-
-@_bool_event_gradient.defjvp
-def _bool_event_gradient_jvp(primals, tangents):
-    t, result = primals
-    tangent_t, _ = tangents
-    out = _bool_event_gradient(t, result)
-    tangent_out = 1.0
-    return out, tangent_out
-
-
 _PRINT_STATIC = False  # used in tests
 
 
@@ -359,14 +344,6 @@ def body_fun_aux(state):
         # everything breaks.) See #143.
         y_error = jtu.tree_map(lambda x: jnp.where(jnp.isnan(x), jnp.inf, x), y_error)
 
-        # Save info for event handling
-        if event is None:
-            tprevprev = None
-            dense_info_for_event = None
-        else:
-            tprevprev = state.tprev
-            dense_info_for_event = dense_info
-
         error_order = solver.error_order(terms)
         (
             keep_step,
@@ -510,6 +487,76 @@ def save_steps(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState:
             )
             dense_save_index = dense_save_index + jnp.where(keep_step, 1, 0)
 
+        if event is None:
+            event_tprev = None
+            event_tnext = None
+            event_dense_info = None
+            event_values = None
+            event_mask = None
+        else:
+            event_tprev = state.tprev
+            event_tnext = state.tnext
+            event_dense_info = dense_info
+            event_values = jtu.tree_map(
+                lambda cond_fn_i: cond_fn_i(
+                    tprev,
+                    y,
+                    args,
+                    terms=terms,
+                    solver=solver,
+                    t0=t0,
+                    t1=t1,
+                    dt0=dt0,
+                    saveat=saveat,
+                    stepsize_controller=stepsize_controller,
+                    max_steps=max_steps,
+                ),
+                event.cond_fn,
+                is_leaf=callable,
+            )
+            had_event = False
+            old_event_values_leaves, old_event_structure = jtu.tree_flatten(
+                state.event_values
+            )
+            new_event_values_leaves, new_event_structure = jtu.tree_flatten(
+                event_values
+            )
+            assert old_event_structure == new_event_structure
+            event_mask_leaves = []
+            for old_event_value_i, new_event_value_i in zip(
+                old_event_values_leaves, new_event_values_leaves
+            ):
+                assert jnp.shape(old_event_value_i) == ()
+                if jnp.shape(new_event_value_i) != ():
+                    raise ValueError(
+                        "Event functions must return a scalar, got shape "
+                        f"{jnp.shape(new_event_value_i)}."
+                    )
+                old_dtype = jnp.result_type(old_event_value_i)
+                new_dtype = jnp.result_type(new_event_value_i)
+                if old_dtype != new_dtype:
+                    raise ValueError(
+                        "Event functions must consistently return either a boolean or "
+                        f"a float, got a change of dtype from {old_dtype} to "
+                        f"{new_dtype}."
+                    )
+                if jnp.issubdtype(new_dtype, jnp.floating):
+                    event_mask_i = jnp.sign(old_event_value_i) != jnp.sign(
+                        new_event_value_i
+                    )
+                elif jnp.issubdtype(new_dtype, jnp.bool_):
+                    event_mask_i = new_event_value_i
+                else:
+                    assert False
+                event_mask_leaves.append(event_mask_i & jnp.invert(had_event))
+                had_event = event_mask_i | had_event
+            event_mask = jtu.tree_unflatten(old_event_structure, event_mask_leaves)
+            result = RESULTS.where(
+                had_event,
+                RESULTS.terminating_event_occurred,
+                state.result,
+            )
+
         new_state = State(
             y=y,
             tprev=tprev,
@@ -522,61 +569,17 @@ def save_steps(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState:
             num_accepted_steps=num_accepted_steps,
             num_rejected_steps=num_rejected_steps,
             save_state=save_state,
-            dense_ts=dense_ts,  # pyright: ignore
+            dense_ts=dense_ts,  # pyright: ignore[reportArgumentType]
             dense_infos=dense_infos,
             dense_save_index=dense_save_index,
             progress_meter_state=progress_meter_state,
-            dense_info_for_event=dense_info_for_event,
-            tprevprev=tprevprev,
-            event_result=None,
-            event_mask=None,
+            event_tprev=event_tprev,
+            event_tnext=event_tnext,
+            event_dense_info=event_dense_info,
+            event_values=event_values,
+            event_mask=event_mask,
         )
 
-        if event is not None:
-
-            def _call_event(_cond_fn):
-                return _cond_fn(
-                    new_state,
-                    y=y,
-                    solver=solver,
-                    stepsize_controller=stepsize_controller,
-                    saveat=saveat,
-                    t0=t0,
-                    t1=t1,
-                    dt0=dt0,
-                    max_steps=max_steps,
-                    terms=terms,
-                    args=args,
-                )
-
-            event_result = jtu.tree_map(_call_event, event.cond_fn, is_leaf=callable)
-            event_mask = jtu.tree_map(_compare_events, state.event_result, event_result)
-            result = RESULTS.where(
-                _event_happened(event_mask),
-                RESULTS.discrete_terminating_event_occurred,
-                result,
-            )
-            # If multiple events are triggered in one step we take the first one
-            event_hist = []
-
-            def counter_update(x):
-                out = jnp.where(jnp.any(jnp.array(event_hist)), False, x)
-                event_hist.append(x)
-                return out
-
-            event_mask = jtu.tree_map(counter_update, event_mask, is_leaf=_is_bool)
-
-            new_state = eqx.tree_at(lambda s: s.result, new_state, result)
-            new_state = eqx.tree_at(
-                lambda s: s.event_result, new_state, event_result, is_leaf=_is_none
-            )
-            new_state = eqx.tree_at(
-                lambda s: s.event_mask,
-                new_state,
-                event_mask,
-                is_leaf=_is_none,
-            )
-
         return (
             new_state,
             (type(new_state.made_jump) is not bool),
@@ -607,16 +610,113 @@ def body_fun(state):
     final_state = outer_while_loop(
         cond_fun, body_fun, init_state, max_steps=max_steps, buffers=_outer_buffers
     )
+    result = final_state.result
+
+    if event is None or event.root_finder is None:
+        tfinal = final_state.tprev
+        yfinal = final_state.y
+    else:
+        # If we're on this branch, it means that an event may have triggered, and now we
+        # may need to do a root find, in order to locate the event time.
+        event_mask = final_state.event_mask
+        flat_mask = jtu.tree_leaves(event_mask)
+        assert all(jnp.shape(x) == () for x in flat_mask)
+        event_happened = jnp.any(jnp.stack(flat_mask))
+
+        def _root_find():
+            _interpolator = solver.interpolation_cls(
+                t0=final_state.event_tprev,
+                t1=final_state.event_tnext,
+                **final_state.event_dense_info,
+            )
+
+            def _to_root_find(_t, _):
+                _distance_from_t_end = final_state.event_tnext - _t
+
+                def _call_real(_event_mask_i, _cond_fn_i):
+                    def _call_real_impl():
+                        # First evaluate the triggered event.
+                        _y = _interpolator.evaluate(_t)
+                        _value = _cond_fn_i(
+                            t=_t,
+                            y=_y,
+                            args=args,
+                            terms=terms,
+                            solver=solver,
+                            t0=t0,
+                            t1=t1,
+                            dt0=dt0,
+                            saveat=saveat,
+                            stepsize_controller=stepsize_controller,
+                            max_steps=max_steps,
+                        )
+                        # Second: if this is a boolean event, then normalise to a
+                        # floating point number by having the root occur at the end of
+                        # the last step, i.e. `event_tnext`.
+                        _value_dtype = jnp.result_type(_value)
+                        if jnp.issubdtype(_value_dtype, jnp.bool_):
+                            _value = _distance_from_t_end
+                        else:
+                            assert jnp.issubdtype(_value_dtype, jnp.floating)
+                        return _value
+
+                    # Only the triggered event actually gets to the decide what time the
+                    # event occurs; everything else is zeroed out to automatically give
+                    # a root.
+                    #
+                    # We allow this `lax.cond` to be inefficiently transformed into a
+                    # `lax.select` when `_event_mask_i` is batched. There isn't any way
+                    # to avoid this, I think.
+                    _value = lax.cond(_event_mask_i, _call_real_impl, lambda: 0.0)
+
+                    # Third: if no events triggered at all, then has the root occur at
+                    # the end of the last step (which will be the `t1` of the overall
+                    # solve).
+                    _value = jnp.where(event_happened, _value, _distance_from_t_end)
+                    return _value
+
+                return jtu.tree_map(
+                    _call_real,
+                    event_mask,
+                    event.cond_fn,
+                )
+
+            _options = {
+                "lower": final_state.event_tprev,
+                "upper": final_state.event_tnext,
+            }
+            _event_root_find = optx.root_find(
+                _to_root_find,
+                event.root_finder,
+                y0=final_state.event_tnext,
+                options=_options,
+                throw=False,
+            )
+            _tfinal = _event_root_find.value
+            # TODO: we might need to change the way we evaluate `_yfinal` in order to
+            # get more accurate derivatives?
+            _yfinal = _interpolator.evaluate(_tfinal)
+            _result = RESULTS.where(
+                _event_root_find.result == optx.RESULTS.successful,
+                result,
+                RESULTS.promote(_event_root_find.result),
+            )
+            return _tfinal, _yfinal, _result
+
+        # Fastpath: if no event happened anywhere at all, then skip the root-find
+        # altogether.
+        # Note that `_root_find` might still be called on batch elements which did not
+        # have an event, so we still need to access `event_happened` inside of it.
+        tfinal, yfinal, result = lax.cond(
+            eqxi.unvmap_any(event_happened),
+            _root_find,
+            lambda: (final_state.tprev, final_state.y, result),
+        )
 
     def _save_t1(subsaveat, save_state):
         if subsaveat.t1 and not subsaveat.steps:
             # If subsaveat.steps then the final value is already saved.
-            #
-            # Use `tprev` instead of `t1` in case of an event terminating the solve
-            # early. (And absent such an event then `tprev == t1`.)
-            save_state = _save(
-                final_state.tprev, final_state.y, args, subsaveat.fn, save_state
-            )
+            save_state = _save(tfinal, yfinal, args, subsaveat.fn, save_state)
         return save_state
 
     save_state = jtu.tree_map(
@@ -627,10 +727,8 @@ def _save_t1(subsaveat, save_state):
     )
 
     final_state = _handle_static(final_state)
-    result = RESULTS.where(
-        cond_fun(final_state), RESULTS.max_steps_reached, final_state.result
-    )
-    aux_stats = dict()
+    result = RESULTS.where(cond_fun(final_state), RESULTS.max_steps_reached, result)
+    aux_stats = dict()  # TODO: put something in here?
     return eqx.tree_at(lambda s: s.result, final_state, result), aux_stats
 
 
@@ -644,6 +742,7 @@ class SaveAt(eqx.Module):  # noqa: F811
 
 
 @eqx.filter_jit
+@eqxi.doc_remove_args("discrete_terminating_event")
 def diffeqsolve(
     terms: PyTree[AbstractTerm],
     solver: AbstractSolver,
@@ -663,6 +762,8 @@ def diffeqsolve(
     solver_state: Optional[PyTree[ArrayLike]] = None,
     controller_state: Optional[PyTree[ArrayLike]] = None,
     made_jump: Optional[BoolScalarLike] = None,
+    # Exists for backward compatibility
+    discrete_terminating_event: Optional[AbstractDiscreteTerminatingEvent] = None,
 ) -> Solution:
     """Solves a differential equation.
 
@@ -707,8 +808,8 @@ def diffeqsolve(
         discretise-then-optimise, which is usually the best option for most problems.
         See the page on [Adjoints](./adjoints.md) for more information.
 
-    - `discrete_terminating_event`: A discrete event at which to terminate the solve
-        early. See the page on [Events](./events.md) for more information.
+    - `event`: An event at which to terminate the solve early. See the page on
+        [Events](./events.md) for more information.
 
     - `max_steps`: The maximum number of steps to take before quitting the computation
         unconditionally.
@@ -718,9 +819,7 @@ def diffeqsolve(
 
     - `throw`: Whether to raise an exception if the integration fails for any reason.
 
-        If `True` then an integration failure will raise an error. Note that the errors
-        are only reliably raised on CPUs. If on GPUs then the error may only be
-        printed to stderr, whilst on TPUs then the behaviour is undefined.
+        If `True` then an integration failure will raise a runtime error.
 
         If `False` then the returned solution object will have a `result` field
         indicating whether any failures occurred.
@@ -771,6 +870,25 @@ def diffeqsolve(
     # Initial set-up
     #
 
+    # Backward compatibility
+    if discrete_terminating_event is not None:
+        warnings.warn(
+            "`diffrax.diffeqsolve(..., discrete_terminating_event=...)` is deprecated "
+            "in favour of the more general `diffrax.diffeqsolve(..., event=...)` "
+            "interface. This will be removed in some future version of Diffrax.",
+            category=DeprecationWarning,
+            stacklevel=2,
+        )
+        if event is None:
+            event = Event(
+                cond_fn=DiscreteTerminatingEventToCondFn(discrete_terminating_event)
+            )
+        else:
+            raise ValueError(
+                "Cannot pass both "
+                "`diffrax.diffeqsolve(..., event=..., discrete_terminating_event=...)`."
+            )
+
     # Error checking
     if dt0 is not None:
         msg = (
@@ -1030,49 +1148,97 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
     num_rejected_steps = 0
     made_jump = False if made_jump is None else made_jump
     result = RESULTS.successful
+    if saveat.dense or event is not None:
+        _, _, dense_info_struct, _, _ = eqx.filter_eval_shape(
+            solver.step, terms, tprev, tnext, y0, args, solver_state, made_jump
+        )
     if saveat.dense:
         if max_steps is None:
             raise ValueError(
                 "`max_steps=None` is incompatible with `saveat.dense=True`"
             )
-        (
-            _,
-            _,
-            dense_info,
-            _,
-            _,
-        ) = eqx.filter_eval_shape(
-            solver.step, terms, tprev, tnext, y0, args, solver_state, made_jump
-        )
         dense_ts = jnp.full(max_steps + 1, jnp.inf, dtype=time_dtype)
         _make_full = lambda x: jnp.full(
             (max_steps,) + jnp.shape(x), jnp.inf, dtype=x.dtype
         )
-        dense_infos = jtu.tree_map(_make_full, dense_info)
+        dense_infos = jtu.tree_map(_make_full, dense_info_struct)  # pyright: ignore[reportPossiblyUnboundVariable]
         dense_save_index = 0
     else:
         dense_ts = None
         dense_infos = None
         dense_save_index = None
 
+    # Progress meter
+    progress_meter_state = progress_meter.init()
+
+    # Events
     if event is None:
-        tprevprev = None
-        dense_info_for_event = None
+        event_tprev = None
+        event_tnext = None
+        event_dense_info = None
+        event_values = None
         event_mask = None
     else:
-        tprevprev = tprev
-        _, _, traced_dense_info_event, _, _ = eqx.filter_eval_shape(
-            solver.step, terms, tprev, tnext, y0, args, solver_state, made_jump
+        event_tprev = tprev
+        event_tnext = tnext
+        # Fill the dense-info with dummy values on the first step, when we haven't yet
+        # made any steps.
+        # Note that we're threading a needle here! What if we terminate on the very
+        # first step? Our dense-info (and thus a subsequent root find) will be
+        # completely wrong!
+        # Fortunately, this can't quite happen:
+        # - A boolean event never uses dense-info (the interpolation is unused and we go
+        #   to the end of the interval).
+        # - A floating event can't terminate on the first step (it requires a sign
+        #   change).
+        event_dense_info = jtu.tree_map(
+            lambda x: jnp.empty(x.shape, x.dtype),
+            dense_info_struct,  # pyright: ignore[reportPossiblyUnboundVariable]
         )
-        dense_info_for_event = jtu.tree_map(
-            lambda x: jnp.empty(x.shape, x.dtype), traced_dense_info_event
-        )
-        event_mask = jtu.tree_map(
-            lambda x: jnp.bool_(False), event.cond_fn, is_leaf=callable
+
+        event_values = jtu.tree_map(
+            lambda cond_fn_i: cond_fn_i(
+                tprev,
+                y0,
+                args,
+                terms=terms,
+                solver=solver,
+                t0=t0,
+                t1=t1,
+                dt0=dt0,
+                saveat=saveat,
+                stepsize_controller=stepsize_controller,
+                max_steps=max_steps,
+            ),
+            event.cond_fn,
+            is_leaf=callable,
         )
 
-    # Progress meter
-    progress_meter_state = progress_meter.init()
+        had_event = False
+        event_values_leaves, event_structure = jtu.tree_flatten(event_values)
+        event_mask_leaves = []
+        for event_value_i in event_values_leaves:
+            if jnp.shape(event_value_i) != ():
+                raise ValueError(
+                    "Event functions must return a scalar, got shape "
+                    f"{jnp.shape(event_value_i)}."
+                )
+            event_dtype = jnp.result_type(event_value_i)
+            if jnp.issubdtype(event_dtype, jnp.floating):
+                event_mask_i = False  # Has not yet had the opportunity to change sign.
+            elif jnp.issubdtype(event_dtype, jnp.bool_):
+                event_mask_i = event_value_i
+            else:
+                assert False
+            event_mask_leaves.append(event_mask_i & jnp.invert(had_event))
+            had_event = event_mask_i | had_event
+        event_mask = jtu.tree_unflatten(event_structure, event_mask_leaves)
+        result = RESULTS.where(
+            had_event,
+            RESULTS.terminating_event_occurred,
+            result,
+        )
+        del had_event, event_values_leaves, event_structure, event_mask_leaves
 
     # Initialise state
     init_state = State(
@@ -1091,35 +1257,13 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
         dense_infos=dense_infos,
         dense_save_index=dense_save_index,
         progress_meter_state=progress_meter_state,
-        tprevprev=tprevprev,
-        dense_info_for_event=dense_info_for_event,
-        event_result=None,
+        event_tprev=event_tprev,
+        event_tnext=event_tnext,
+        event_dense_info=event_dense_info,
+        event_values=event_values,
         event_mask=event_mask,
     )
 
-    # Since cond_fn might depend on the state, we need to initialise event_result here.
-    if event is not None:
-
-        def _call_event(_cond_fn):
-            return _cond_fn(
-                init_state,
-                y=y0,
-                solver=solver,
-                stepsize_controller=stepsize_controller,
-                saveat=saveat,
-                t0=t0,
-                t1=t1,
-                dt0=dt0,
-                max_steps=max_steps,
-                terms=terms,
-                args=args,
-            )
-
-        event_result = jtu.tree_map(_call_event, event.cond_fn, is_leaf=callable)
-        init_state = eqx.tree_at(
-            lambda s: s.event_result, init_state, event_result, is_leaf=_is_none
-        )
-
     #
     # Main loop
     #
@@ -1153,75 +1297,9 @@ def _call_event(_cond_fn):
     )
     ys = jtu.tree_map(lambda s: s.ys, final_state.save_state, is_leaf=is_save_state)
 
-    # Do root find for exact event times
-    if event is not None:
-        event_mask = final_state.event_mask
-        event_happened = _event_happened(event_mask)
-        interpolator = solver.interpolation_cls(
-            t0=final_state.tprevprev,
-            t1=final_state.tprev,
-            **final_state.dense_info_for_event,
-        )
-        if event.root_finder is None:
-            tevent = final_state.tprev
-        else:
-
-            def _to_root_find(t, args):
-                def _call_real(cond_fn_i, event_mask_i):
-                    y = interpolator.evaluate(t)
-                    result = cond_fn_i(
-                        final_state,
-                        y=y,
-                        solver=solver,
-                        stepsize_controller=stepsize_controller,
-                        saveat=saveat,
-                        t0=t0,
-                        t1=t1,
-                        dt0=dt0,
-                        max_steps=max_steps,
-                        terms=terms,
-                        args=args,
-                    )
-
-                    if jnp.result_type(result) == jnp.bool_:
-                        result = final_state.tprev - t
-                    return jnp.where(event_mask_i, result, 0.0)
-
-                results = jtu.tree_map(
-                    _call_real,
-                    event.cond_fn,
-                    event_mask,
-                    is_leaf=callable,
-                )
-                # If no events are triggered simply push tevent towards tprev
-                results = jtu.tree_map(
-                    lambda x: jnp.where(event_happened, x, final_state.tprev - t),
-                    results,
-                )
-                return results
-
-            options = {"lower": final_state.tprevprev, "upper": final_state.tprev}
-            roots = optx.root_find(
-                _to_root_find, event.root_finder, y0=final_state.tprev, options=options
-            )
-            tevent = roots.value
-
-        # We might need to change this in order to get more accurate derivatives
-        yevent = interpolator.evaluate(tevent)
-
-        def _save_yevent(subsaveat: SubSaveAt, y):
-            return jtu.tree_map(
-                lambda s, _y: _y.at[-1].set(s), subsaveat.fn(tevent, yevent, args), y
-            )
-
-        ys = jtu.tree_map(_save_yevent, saveat.subs, ys, is_leaf=_is_subsaveat)
-        ts = jtu.tree_map(lambda _t: _t.at[-1].set(tevent), ts)
-    else:
-        event_mask = None
-
     # It's important that we don't do any further postprocessing on `ys` here, as
     # it is the `final_state` value that is used when backpropagating via
-    # optimise-then-discretise.
+    # `BacksolveAdjoint`.
 
     if saveat.controller_state:
         controller_state = final_state.controller_state
@@ -1257,8 +1335,10 @@ def _save_yevent(subsaveat: SubSaveAt, y):
         "num_accepted_steps": final_state.num_accepted_steps,
         "num_rejected_steps": final_state.num_rejected_steps,
         "max_steps": max_steps,
+        **aux_stats,
     }
     result = final_state.result
+    event_mask = final_state.event_mask
     sol = Solution(
         t0=t0,
         t1=t1,
diff --git a/diffrax/_solution.py b/diffrax/_solution.py
index 35fcd51e..f79b60c5 100644
--- a/diffrax/_solution.py
+++ b/diffrax/_solution.py
@@ -1,3 +1,4 @@
+import warnings
 from typing import Any, Optional
 
 import jax
@@ -11,16 +12,32 @@
 
 class RESULTS(optx.RESULTS):  # pyright: ignore
     successful = ""
-    discrete_terminating_event_occurred = (
-        "Terminating differential equation solve because a discrete terminating event "
-        "occurred."
-    )
     max_steps_reached = (
         "The maximum number of solver steps was reached. Try increasing `max_steps`."
     )
     dt_min_reached = (
         "The minimum step size was reached in the differential equation solver."
     )
+    terminating_event_occurred = (
+        "Terminating differential equation solve because an event occurred."
+    )
+
+
+# Backward compatibility
+# Evil monkey-patching so that we don't mess with how `Enumeration`s work.
+
+
+@property
+def discrete_terminating_event_occurred(self):
+    warnings.warn(
+        "`diffrax.RESULTS.discrete_terminating_event_occurred` is deprecated in "
+        "favour of `diffrax.RESULTS.terminating_event_occurred`. This will be "
+        "removed in some future version of Diffrax."
+    )
+    return self.terminating_event_occurred
+
+
+RESULTS.discrete_terminating_event_occurred = discrete_terminating_event_occurred  # pyright: ignore[reportAttributeAccessIssue]
 
 
 def is_okay(result: RESULTS) -> Bool[Array, ""]:
@@ -35,7 +52,7 @@ def is_successful(result: RESULTS) -> Bool[Array, ""]:
 # TODO: In the future we may support other event types, in which case this function
 # should be updated.
 def is_event(result: RESULTS) -> Bool[Array, ""]:
-    return result == RESULTS.discrete_terminating_event_occurred
+    return result == RESULTS.terminating_event_occurred
 
 
 def update_result(old_result: RESULTS, new_result: RESULTS) -> RESULTS:
diff --git a/docs/api/events.md b/docs/api/events.md
index 263a195c..fa0fe75e 100644
--- a/docs/api/events.md
+++ b/docs/api/events.md
@@ -1,24 +1,10 @@
 # Events
 
-Events allow for interrupting a differential equation solve, and changing its internal state, or terminating the solve before `t1` is reached.
+Events allow for interrupting a differential equation solve, by terminating the solve before `t1` is reached.
 
-At the moment a single kind of event is supported: discrete events which are checked at the end of every step, and which halt the integration once they become true.
-
-??? abstract "`diffrax.AbstractDiscreteTerminatingEvent`"
-
-    ::: diffrax.AbstractDiscreteTerminatingEvent
-        selection:
-            members:
-                - __call__
-
----
-
-::: diffrax.DiscreteTerminatingEvent
+::: diffrax.Event
     selection:
         members:
             - __init__
 
-::: diffrax.SteadyStateEvent
-    selection:
-        members:
-            - __init__
+::: diffrax.steady_state_event
diff --git a/test/test_adjoint.py b/test/test_adjoint.py
index f0a15e71..834b7a3b 100644
--- a/test/test_adjoint.py
+++ b/test/test_adjoint.py
@@ -258,13 +258,13 @@ def run(model):
     run(mlp)
 
 
-@pytest.mark.skip(reason="SteadyStateEvent discontinued")
 def test_implicit():
     class ExponentialDecayToSteadyState(eqx.Module):
         steady_state: Array
         non_jax_type: Any
 
         def __call__(self, t, y, args):
+            del t, args
             return self.steady_state - y
 
     def loss(model, target_steady_state):
@@ -276,7 +276,7 @@ def loss(model, target_steady_state):
         y0 = 1.0
         max_steps = None
         controller = diffrax.PIDController(rtol=1e-3, atol=1e-6)
-        event = diffrax.SteadyStateEvent()  # pyright: ignore
+        event = diffrax.SteadyStateEvent()
         adjoint = diffrax.ImplicitAdjoint()
         sol = diffrax.diffeqsolve(
             term,
@@ -287,7 +287,7 @@ def loss(model, target_steady_state):
             y0,
             max_steps=max_steps,
             stepsize_controller=controller,
-            discrete_terminating_event=event,  # pyright: ignore
+            discrete_terminating_event=event,
             adjoint=adjoint,
         )
         (y1,) = cast(Array, sol.ys)
diff --git a/test/test_event.py b/test/test_event.py
index d9334e08..705446ad 100644
--- a/test/test_event.py
+++ b/test/test_event.py
@@ -8,7 +8,6 @@
 from jaxtyping import Array
 
 
-@pytest.mark.skip(reason="Old event implementation")
 def test_discrete_terminate1():
     term = diffrax.ODETerm(lambda t, y, args: y)
     solver = diffrax.Tsit5()
@@ -18,23 +17,24 @@ def test_discrete_terminate1():
     y0 = 1.0
 
     def event_fn(state, **kwargs):
+        del kwargs
         assert isinstance(state.y, jax.Array)
         return state.tprev > 10
 
-    event = diffrax.DiscreteTerminatingEvent(event_fn)  # pyright: ignore
-    sol = diffrax.diffeqsolve(
-        term,
-        solver,
-        t0,
-        t1,
-        dt0,
-        y0,
-        discrete_terminating_event=event,  # pyright: ignore
-    )
+    event = diffrax.DiscreteTerminatingEvent(event_fn)
+    with pytest.warns(DeprecationWarning, match="discrete_terminating_event"):
+        sol = diffrax.diffeqsolve(
+            term,
+            solver,
+            t0,
+            t1,
+            dt0,
+            y0,
+            discrete_terminating_event=event,
+        )
     assert jnp.all(cast(Array, sol.ys) > 10)
 
 
-@pytest.mark.skip(reason="Old event implementation")
 def test_discrete_terminate2():
     term = diffrax.ODETerm(lambda t, y, args: y)
     solver = diffrax.Tsit5()
@@ -44,23 +44,24 @@ def test_discrete_terminate2():
     y0 = 1.0
 
     def event_fn(state, **kwargs):
+        del kwargs
         assert isinstance(state.y, jax.Array)
         return state.tprev > 10
 
-    event = diffrax.DiscreteTerminatingEvent(event_fn)  # pyright: ignore
-    sol = diffrax.diffeqsolve(
-        term,
-        solver,
-        t0,
-        t1,
-        dt0,
-        y0,
-        discrete_terminating_event=event,  # pyright: ignore
-    )
+    event = diffrax.DiscreteTerminatingEvent(event_fn)
+    with pytest.warns(DeprecationWarning, match="discrete_terminating_event"):
+        sol = diffrax.diffeqsolve(
+            term,
+            solver,
+            t0,
+            t1,
+            dt0,
+            y0,
+            discrete_terminating_event=event,
+        )
     assert jnp.all(cast(Array, sol.ts) > 10)
 
 
-@pytest.mark.skip(reason="Old event implementation")
 def test_event_backsolve():
     term = diffrax.ODETerm(lambda t, y, args: y)
     solver = diffrax.Tsit5()
@@ -70,24 +71,26 @@ def test_event_backsolve():
     y0 = 1.0
 
     def event_fn(state, **kwargs):
+        del kwargs
         assert isinstance(state.y, jax.Array)
         return state.tprev > 10
 
-    event = diffrax.DiscreteTerminatingEvent(event_fn)  # pyright: ignore
+    event = diffrax.DiscreteTerminatingEvent(event_fn)
 
     @jax.jit
     @jax.grad
     def run(y0):
-        sol = diffrax.diffeqsolve(
-            term,
-            solver,
-            t0,
-            t1,
-            dt0,
-            y0,
-            discrete_terminating_event=event,  # pyright: ignore
-            adjoint=diffrax.BacksolveAdjoint(),
-        )
+        with pytest.warns(DeprecationWarning, match="discrete_terminating_event"):
+            sol = diffrax.diffeqsolve(
+                term,
+                solver,
+                t0,
+                t1,
+                dt0,
+                y0,
+                discrete_terminating_event=event,
+                adjoint=diffrax.BacksolveAdjoint(),
+            )
         return jnp.sum(cast(Array, sol.ys))
 
     # And in particular not some other error.
@@ -106,9 +109,10 @@ def test_continuous_terminate1():
     dt0 = 1
     y0 = 1.0
 
-    def cond_fn(state, **kwargs):
-        assert isinstance(state.y, jax.Array)
-        return state.tprev > 10
+    def cond_fn(t, y, args, **kwargs):
+        del args, kwargs
+        assert isinstance(y, jax.Array)
+        return t > 10
 
     event = diffrax.Event(cond_fn=cond_fn)
     sol = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
@@ -123,9 +127,10 @@ def test_continuous_terminate2():
     dt0 = 1
     y0 = 1.0
 
-    def cond_fn(state, **kwargs):
-        assert isinstance(state.y, jax.Array)
-        return state.tprev - 10.0
+    def cond_fn(t, y, args, **kwargs):
+        del args, kwargs
+        assert isinstance(y, jax.Array)
+        return t - 10.0
 
     event = diffrax.Event(cond_fn=cond_fn)
     sol = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
@@ -140,8 +145,9 @@ def test_continuous_event_time():
     dt0 = 1.0
     y0 = 1.0
 
-    def cond_fn(state, y, **kwargs):
-        assert isinstance(state.y, jax.Array)
+    def cond_fn(t, y, args, **kwargs):
+        del t, args, kwargs
+        assert isinstance(y, jax.Array)
         return y - jnp.exp(1.0)
 
     root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
@@ -158,8 +164,9 @@ def test_continuous_event_value():
     dt0 = 1.0
     y0 = -10.0
 
-    def cond_fn(state, y, **kwargs):
-        assert isinstance(state.y, jax.Array)
+    def cond_fn(t, y, args, **kwargs):
+        del t, args, kwargs
+        assert isinstance(y, jax.Array)
         return y
 
     root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
@@ -176,8 +183,9 @@ def test_continuous_no_event():
     dt0 = 1.0
     y0 = -10.0
 
-    def cond_fn(state, y, **kwargs):
-        assert isinstance(state.y, jax.Array)
+    def cond_fn(t, y, args, **kwargs):
+        del t, args, kwargs
+        assert isinstance(y, jax.Array)
         return y
 
     root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
@@ -195,12 +203,14 @@ def test_continuous_two_events():
     dt0 = 1.0
     y0 = -10.0
 
-    def cond_fn_1(state, y, **kwargs):
-        assert isinstance(state.y, jax.Array)
+    def cond_fn_1(t, y, args, **kwargs):
+        del t, args, kwargs
+        assert isinstance(y, jax.Array)
         return y
 
-    def cond_fn_2(state, y, **kwargs):
-        assert isinstance(state.y, jax.Array)
+    def cond_fn_2(t, y, args, **kwargs):
+        del t, args, kwargs
+        assert isinstance(y, jax.Array)
         return y + 5.0
 
     root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
@@ -213,12 +223,14 @@ def cond_fn_2(state, y, **kwargs):
 
 def test_continuous_event_time_grad():
     def vector_field(t, y, args):
-        x, v = y
+        del t, args
+        _, v = y
         d_out = v, -8.0
         return jnp.array(d_out)
 
-    def cond_fn(state, y, **kwargs):
-        x, v = y
+    def cond_fn(t, y, args, **kwargs):
+        del t, args, kwargs
+        x, _ = y
         return x
 
     term = diffrax.ODETerm(vector_field)
@@ -250,7 +262,7 @@ def phi(_x):
             return _sol.ys[-1, :]
 
         def event_fn(_y):
-            return cond_fn(None, _y)
+            return cond_fn(t=None, y=_y, args=None)
 
         _, num = jax.vjp(phi, x0)
         (num,) = num(jax.grad(event_fn)(y0))

From e2ab3ce240f67b472bb8d5daf1a03e703078c3a6 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Thu, 23 May 2024 17:31:22 +0200
Subject: [PATCH 06/28] Save values returned by root find when

---
 diffrax/_integrate.py | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index 0b58e7c7..fa2dfc23 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -712,11 +712,22 @@ def _call_real_impl():
             _root_find,
             lambda: (final_state.tprev, final_state.y, result),
         )
+        # Update save_index to replace last saved step with event values
+        save_index = final_state.save_state.save_index - 1
+        final_state = eqx.tree_at(
+            lambda s: s.save_state.save_index, final_state, save_index
+        )
 
     def _save_t1(subsaveat, save_state):
-        if subsaveat.t1 and not subsaveat.steps:
-            # If subsaveat.steps then the final value is already saved.
-            save_state = _save(tfinal, yfinal, args, subsaveat.fn, save_state)
+        if event.root_finder is None:
+            if subsaveat.t1 and not subsaveat.steps:
+                # If subsaveat.steps then the final value is already saved.
+                save_state = _save(tfinal, yfinal, args, subsaveat.fn, save_state)
+        else:
+            if subsaveat.t1 or subsaveat.steps:
+                # In this branch we need to replace the last value with tfinal
+                # and yfinal returned by the root finder also if subsaveat.steps.
+                save_state = _save(tfinal, yfinal, args, subsaveat.fn, save_state)
         return save_state
 
     save_state = jtu.tree_map(

From 0853e49e0c017100db9d948e1a9a014e238e3c5b Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Fri, 24 May 2024 11:35:30 +0200
Subject: [PATCH 07/28]  now returns condition function

---
 diffrax/_event.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/diffrax/_event.py b/diffrax/_event.py
index 6f3cbc58..54c1ef9f 100644
--- a/diffrax/_event.py
+++ b/diffrax/_event.py
@@ -71,7 +71,7 @@ def _cond_fn(t, y, args, *, terms, solver, stepsize_controller, **kwargs):
         vf = solver.func(terms, t, y, args)
         return _norm(vf) < _atol + _rtol * _norm(y)
 
-    return Event(cond_fn=_cond_fn)
+    return _cond_fn
 
 
 #
@@ -99,7 +99,7 @@ class SteadyStateEvent(AbstractDiscreteTerminatingEvent):
     norm: Callable[[PyTree[Array]], RealScalarLike] = optx.rms_norm
 
     def __call__(self, state, *, args, **kwargs):
-        return steady_state_event(self.rtol, self.atol, self.norm).cond_fn(
+        return steady_state_event(self.rtol, self.atol, self.norm)(
             state.tprev, state.y, args, **kwargs
         )
 

From 1488206382d32a130d3583dcb9818fba3b0e40ba Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Fri, 24 May 2024 13:30:18 +0200
Subject: [PATCH 08/28] Fixed error for . All tests pass now.

---
 diffrax/_integrate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index fa2dfc23..85798ca8 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -719,7 +719,7 @@ def _call_real_impl():
         )
 
     def _save_t1(subsaveat, save_state):
-        if event.root_finder is None:
+        if event is None or event.root_finder is None:
             if subsaveat.t1 and not subsaveat.steps:
                 # If subsaveat.steps then the final value is already saved.
                 save_state = _save(tfinal, yfinal, args, subsaveat.fn, save_state)

From 4872009d8534f4a6efe86fdcd5c9815a3c6e1c74 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Sun, 26 May 2024 17:24:55 +0200
Subject: [PATCH 09/28] Added additional tests

Added a bunch of additional tests for events.

Also changed the way `save_index` was updated to handle PyTrees of subsaveats.
---
 diffrax/_integrate.py |   7 +-
 test/test_event.py    | 249 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 250 insertions(+), 6 deletions(-)

diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index 85798ca8..d22d80b5 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -713,7 +713,12 @@ def _call_real_impl():
             lambda: (final_state.tprev, final_state.y, result),
         )
         # Update save_index to replace last saved step with event values
-        save_index = final_state.save_state.save_index - 1
+        save_index = jtu.tree_map(
+            lambda _, s: s.save_index - 1,
+            saveat.subs,
+            save_state,
+            is_leaf=_is_subsaveat,
+        )
         final_state = eqx.tree_at(
             lambda s: s.save_state.save_index, final_state, save_index
         )
diff --git a/test/test_event.py b/test/test_event.py
index 705446ad..4a323e98 100644
--- a/test/test_event.py
+++ b/test/test_event.py
@@ -1,3 +1,4 @@
+import functools as ft
 from typing import cast
 
 import diffrax
@@ -101,6 +102,46 @@ def run(y0):
 # diffrax.SteadyStateEvent tested as part of test_adjoint.py::test_implicit
 
 
+def test_steady_state_event():
+    term = diffrax.ODETerm(lambda t, y, args: -1.0 * y)
+    controller = diffrax.PIDController(rtol=1e-3, atol=1e-6)
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = jnp.inf
+    dt0 = 1
+    y0 = 1.0
+    cond_fn = diffrax.steady_state_event()
+    event = diffrax.Event(cond_fn)
+    sol = diffrax.diffeqsolve(
+        term, solver, t0, t1, dt0, y0, stepsize_controller=controller, event=event
+    )
+
+    assert cast(Array, sol.event_mask)
+    assert jnp.all(jnp.isclose(cast(Array, sol.ys), 0.0, atol=1e-5))
+
+
+def test_no_step_event():
+    term = diffrax.ODETerm(lambda t, y, args: jnp.array([1, 1]))
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = 10
+    dt0 = 1
+    y0 = jnp.array([1, -1e-1])
+
+    def cond_fn(t, y, args, **kwargs):
+        del t, args, kwargs
+        assert isinstance(y, Array)
+        _, x = y
+        return x < 0
+
+    event = diffrax.Event(cond_fn)
+    sol = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
+
+    assert sol.stats["num_steps"] == 0
+    assert jnp.all(jnp.isclose(cast(Array, sol.ys)[-1], y0, 1e-7))
+    assert jnp.all(jnp.isclose(cast(Array, sol.ts), t0, 1e-7))
+
+
 def test_continuous_terminate1():
     term = diffrax.ODETerm(lambda t, y, args: y)
     solver = diffrax.Tsit5()
@@ -110,8 +151,7 @@ def test_continuous_terminate1():
     y0 = 1.0
 
     def cond_fn(t, y, args, **kwargs):
-        del args, kwargs
-        assert isinstance(y, jax.Array)
+        del y, args, kwargs
         return t > 10
 
     event = diffrax.Event(cond_fn=cond_fn)
@@ -204,9 +244,9 @@ def test_continuous_two_events():
     y0 = -10.0
 
     def cond_fn_1(t, y, args, **kwargs):
-        del t, args, kwargs
-        assert isinstance(y, jax.Array)
-        return y
+        del y, args, kwargs
+        assert isinstance(t, jax.Array)
+        return t - 10
 
     def cond_fn_2(t, y, args, **kwargs):
         del t, args, kwargs
@@ -273,3 +313,202 @@ def event_fn(_y):
     x0_autograd = jax.vmap(first_bounce_time)(x0_test)
     x0_truegrad = jax.vmap(first_bounce_time_grad)(x0_test)
     assert jnp.all(jnp.isclose(x0_autograd, x0_truegrad, 1e-5))
+
+
+def test_adaptive_stepping_event():
+    term = diffrax.ODETerm(lambda t, y, args: -y)
+    controller = diffrax.PIDController(rtol=1e-3, atol=1e-6)
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = jnp.inf
+    dt0 = 1
+
+    def cond_fn(t, y, args, **kwargs):
+        del t, args, kwargs
+        assert isinstance(y, Array)
+        return y - 1
+
+    root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
+    event = diffrax.Event(cond_fn, root_finder)
+
+    @jax.jit
+    @jax.value_and_grad
+    def run(y):
+        sol = diffrax.diffeqsolve(
+            term, solver, t0, t1, dt0, y, stepsize_controller=controller, event=event
+        )
+        return cast(Array, sol.ys)[-1]
+
+    y0s = [10.0, 100.0, 1000.0]
+    for y0 in y0s:
+        val, grad = run(y0)
+        assert jnp.all(jnp.isclose(val - 1, 0.0, atol=1e-5))
+        assert not jnp.isnan(grad).any()
+
+
+@pytest.mark.parametrize(
+    "stepsize_controller",
+    (diffrax.ConstantStepSize(), diffrax.PIDController(rtol=1e-3, atol=1e-6)),
+)
+def test_event_vmap_y0(stepsize_controller):
+    term = diffrax.ODETerm(lambda t, y, args: -y)
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = jnp.inf
+    dt0 = 1
+
+    def cond_fn(t, y, args, **kwargs):
+        del t, args, kwargs
+        assert isinstance(y, Array)
+        return y - 1
+
+    root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
+    event = diffrax.Event(cond_fn, root_finder)
+
+    @jax.vmap
+    @jax.value_and_grad
+    def run(y):
+        sol = diffrax.diffeqsolve(
+            term,
+            solver,
+            t0,
+            t1,
+            dt0,
+            y,
+            stepsize_controller=stepsize_controller,
+            event=event,
+        )
+        return cast(Array, sol.ys)[-1]
+
+    y0s = jnp.arange(10.0) + 2.0
+    vals, grads = run(y0s)
+    for val, grad in zip(vals, grads):
+        assert jnp.all(jnp.isclose(val - 1, 0.0, atol=1e-5))
+        assert not jnp.isnan(grad).any()
+
+
+@pytest.mark.parametrize(
+    "stepsize_controller",
+    (diffrax.ConstantStepSize(), diffrax.PIDController(rtol=1e-3, atol=1e-6)),
+)
+def test_event_vmap_t0(stepsize_controller):
+    term = diffrax.ODETerm(lambda t, y, args: -y)
+    solver = diffrax.Tsit5()
+    t1 = jnp.inf
+    dt0 = 1
+    y0 = 10
+
+    def cond_fn(t, y, args, **kwargs):
+        del t, args, kwargs
+        assert isinstance(y, Array)
+        return y - 1
+
+    root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
+    event = diffrax.Event(cond_fn, root_finder)
+
+    @jax.vmap
+    @jax.value_and_grad
+    def run(t):
+        sol = diffrax.diffeqsolve(
+            term,
+            solver,
+            t,
+            t1,
+            dt0,
+            y0,
+            stepsize_controller=stepsize_controller,
+            event=event,
+        )
+        return cast(Array, sol.ys)[-1]
+
+    t0s = jnp.arange(10.0) / 10
+    vals, grads = run(t0s)
+    for val, grad in zip(vals, grads):
+        assert jnp.all(jnp.isclose(val - 1, 0.0, atol=1e-5))
+        assert not jnp.isnan(grad).any()
+
+
+@pytest.mark.parametrize(
+    "stepsize_controller",
+    (diffrax.ConstantStepSize(), diffrax.PIDController(rtol=1e-3, atol=1e-6)),
+)
+def test_event_vmap_event_def(stepsize_controller):
+    term = diffrax.ODETerm(lambda t, y, args: -y)
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = jnp.inf
+    dt0 = 1
+    y0 = 10
+
+    def cond_fn(thr, t, y, args, **kwargs):
+        del t, args, kwargs
+        assert isinstance(y, Array)
+        return y - thr
+
+    root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
+
+    @jax.vmap
+    @jax.value_and_grad
+    def run(thr):
+        _cond_fn = ft.partial(cond_fn, thr)
+        event = diffrax.Event(_cond_fn, root_finder)
+        sol = diffrax.diffeqsolve(
+            term,
+            solver,
+            t0,
+            t1,
+            dt0,
+            y0,
+            stepsize_controller=stepsize_controller,
+            event=event,
+        )
+        return cast(Array, sol.ys)[-1]
+
+    thrs = (jnp.arange(10.0) + 1) / 5
+    vals, grads = run(thrs)
+    for thr, val, grad in zip(thrs, vals, grads):
+        assert jnp.all(jnp.isclose(val - thr, 0.0, atol=1e-5))
+        assert not jnp.isnan(grad).any()
+
+
+@pytest.mark.parametrize(
+    "stepsize_controller",
+    (diffrax.ConstantStepSize(), diffrax.PIDController(rtol=1e-3, atol=1e-6)),
+)
+def test_event_vmap_cond_fn(stepsize_controller):
+    term = diffrax.ODETerm(lambda t, y, args: -y)
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = jnp.inf
+    dt0 = 1
+
+    def cond_fn(t, y, args, **kwargs):
+        del t, args, kwargs
+
+        @jax.vmap
+        def _cond_fn(y):
+            return y - 1
+
+        return jnp.max(_cond_fn(y))
+
+    root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
+    event = diffrax.Event(cond_fn, root_finder)
+
+    @jax.value_and_grad
+    def run(y):
+        sol = diffrax.diffeqsolve(
+            term,
+            solver,
+            t0,
+            t1,
+            dt0,
+            y,
+            stepsize_controller=stepsize_controller,
+            event=event,
+        )
+        return jnp.max(cast(Array, sol.ys)[-1])
+
+    y0s = jnp.arange(10.0) + 2.0
+    val, grad = run(y0s)
+    assert jnp.all(jnp.isclose(val - 1, 0.0, atol=1e-5))
+    assert not jnp.isnan(grad).any()

From a1f577c53201c3f813eb0879f96c81d3ed91f507 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Mon, 27 May 2024 12:59:40 +0200
Subject: [PATCH 10/28] Fixed save_index update and shape+dtype check for
 cond_fn

---
 diffrax/_integrate.py | 92 ++++++++++++++++++++++++++-----------------
 1 file changed, 55 insertions(+), 37 deletions(-)

diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index d22d80b5..22d336c5 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -497,8 +497,9 @@ def save_steps(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState:
             event_tprev = state.tprev
             event_tnext = state.tnext
             event_dense_info = dense_info
-            event_values = jtu.tree_map(
-                lambda cond_fn_i: cond_fn_i(
+
+            def _outer_cond_fn(cond_fn_i, old_event_value_i):
+                new_event_value_i = cond_fn_i(
                     tprev,
                     y,
                     args,
@@ -510,22 +511,7 @@ def save_steps(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState:
                     saveat=saveat,
                     stepsize_controller=stepsize_controller,
                     max_steps=max_steps,
-                ),
-                event.cond_fn,
-                is_leaf=callable,
-            )
-            had_event = False
-            old_event_values_leaves, old_event_structure = jtu.tree_flatten(
-                state.event_values
-            )
-            new_event_values_leaves, new_event_structure = jtu.tree_flatten(
-                event_values
-            )
-            assert old_event_structure == new_event_structure
-            event_mask_leaves = []
-            for old_event_value_i, new_event_value_i in zip(
-                old_event_values_leaves, new_event_values_leaves
-            ):
+                )
                 assert jnp.shape(old_event_value_i) == ()
                 if jnp.shape(new_event_value_i) != ():
                     raise ValueError(
@@ -548,13 +534,30 @@ def save_steps(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState:
                     event_mask_i = new_event_value_i
                 else:
                     assert False
+                return new_event_value_i, event_mask_i
+
+            event_values__mask = jtu.tree_map(
+                _outer_cond_fn,
+                event.cond_fn,
+                state.event_values,
+                is_leaf=callable,
+            )
+            event_structure = jtu.tree_structure(event.cond_fn, is_leaf=callable)
+            event_values, event_mask = jtu.tree_transpose(
+                event_structure,
+                None,  # pyright: ignore
+                event_values__mask,
+            )
+            had_event = False
+            event_mask_leaves = []
+            for event_mask_i in jtu.tree_leaves(event_mask):
                 event_mask_leaves.append(event_mask_i & jnp.invert(had_event))
                 had_event = event_mask_i | had_event
-            event_mask = jtu.tree_unflatten(old_event_structure, event_mask_leaves)
+            event_mask = jtu.tree_unflatten(event_structure, event_mask_leaves)
             result = RESULTS.where(
                 had_event,
                 RESULTS.terminating_event_occurred,
-                state.result,
+                result,
             )
 
         new_state = State(
@@ -712,15 +715,22 @@ def _call_real_impl():
             _root_find,
             lambda: (final_state.tprev, final_state.y, result),
         )
+
         # Update save_index to replace last saved step with event values
-        save_index = jtu.tree_map(
-            lambda _, s: s.save_index - 1,
-            saveat.subs,
+        def _update_index(save_state):
+            return eqx.tree_at(
+                lambda s: s.save_index, save_state, replace_fn=lambda i: i - 1
+            )
+
+        save_state = jtu.tree_map(
+            lambda s: _update_index(s),
             save_state,
-            is_leaf=_is_subsaveat,
+            is_leaf=lambda s: isinstance(s, SaveState),
         )
         final_state = eqx.tree_at(
-            lambda s: s.save_state.save_index, final_state, save_index
+            lambda s: s.save_state,
+            final_state,
+            save_state,
         )
 
     def _save_t1(subsaveat, save_state):
@@ -1212,8 +1222,8 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
             dense_info_struct,  # pyright: ignore[reportPossiblyUnboundVariable]
         )
 
-        event_values = jtu.tree_map(
-            lambda cond_fn_i: cond_fn_i(
+        def _outer_cond_fn(cond_fn_i):
+            event_value_i = cond_fn_i(
                 tprev,
                 y0,
                 args,
@@ -1225,15 +1235,7 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
                 saveat=saveat,
                 stepsize_controller=stepsize_controller,
                 max_steps=max_steps,
-            ),
-            event.cond_fn,
-            is_leaf=callable,
-        )
-
-        had_event = False
-        event_values_leaves, event_structure = jtu.tree_flatten(event_values)
-        event_mask_leaves = []
-        for event_value_i in event_values_leaves:
+            )
             if jnp.shape(event_value_i) != ():
                 raise ValueError(
                     "Event functions must return a scalar, got shape "
@@ -1246,6 +1248,22 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
                 event_mask_i = event_value_i
             else:
                 assert False
+            return event_value_i, event_mask_i
+
+        event_values__mask = jtu.tree_map(
+            _outer_cond_fn,
+            event.cond_fn,
+            is_leaf=callable,
+        )
+        event_structure = jtu.tree_structure(event.cond_fn, is_leaf=callable)
+        event_values, event_mask = jtu.tree_transpose(
+            event_structure,
+            None,  # pyright: ignore
+            event_values__mask,
+        )
+        had_event = False
+        event_mask_leaves = []
+        for event_mask_i in jtu.tree_leaves(event_mask):
             event_mask_leaves.append(event_mask_i & jnp.invert(had_event))
             had_event = event_mask_i | had_event
         event_mask = jtu.tree_unflatten(event_structure, event_mask_leaves)
@@ -1254,7 +1272,7 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
             RESULTS.terminating_event_occurred,
             result,
         )
-        del had_event, event_values_leaves, event_structure, event_mask_leaves
+        del had_event, event_structure, event_mask_leaves, event_values__mask
 
     # Initialise state
     init_state = State(

From 95ac30f4e9ce9907b365c440e461a93b8cc37a3d Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Mon, 27 May 2024 13:15:30 +0200
Subject: [PATCH 11/28] Added PyTree check in _outer_cond_fn

---
 diffrax/_integrate.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index 22d336c5..e29f00cf 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -513,6 +513,11 @@ def _outer_cond_fn(cond_fn_i, old_event_value_i):
                     max_steps=max_steps,
                 )
                 assert jnp.shape(old_event_value_i) == ()
+                if jtu.tree_structure(new_event_value_i) != jtu.tree_structure(0):
+                    raise ValueError(
+                        "Event functions must return a scalar, got PyTree with shape "
+                        f"{jtu.tree_structure(new_event_value_i)}."
+                    )
                 if jnp.shape(new_event_value_i) != ():
                     raise ValueError(
                         "Event functions must return a scalar, got shape "
@@ -1236,6 +1241,11 @@ def _outer_cond_fn(cond_fn_i):
                 stepsize_controller=stepsize_controller,
                 max_steps=max_steps,
             )
+            if jtu.tree_structure(event_value_i) != jtu.tree_structure(0):
+                raise ValueError(
+                    "Event functions must return a scalar, got PyTree with shape "
+                    f"{jtu.tree_structure(event_value_i)}."
+                )
             if jnp.shape(event_value_i) != ():
                 raise ValueError(
                     "Event functions must return a scalar, got shape "

From 0c820f3a3ff1d2adccf724a88865c804301a1322 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Mon, 27 May 2024 13:22:50 +0200
Subject: [PATCH 12/28] Added tests for checking that events error out
 correctly under misspecified cond_fn

---
 test/test_event.py | 48 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 48 insertions(+)

diff --git a/test/test_event.py b/test/test_event.py
index 4a323e98..6214c575 100644
--- a/test/test_event.py
+++ b/test/test_event.py
@@ -512,3 +512,51 @@ def run(y):
     val, grad = run(y0s)
     assert jnp.all(jnp.isclose(val - 1, 0.0, atol=1e-5))
     assert not jnp.isnan(grad).any()
+
+
+def test_event_scalar_error():
+    term = diffrax.ODETerm(lambda t, y, args: y)
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = jnp.inf
+    dt0 = 1
+    y0 = 1.0
+
+    def cond_fn_1(t, y, args, **kwargs):
+        del t, args, kwargs
+        assert isinstance(y, jax.Array)
+        return (y,)
+
+    def cond_fn_2(t, y, args, **kwargs):
+        del t, args, kwargs
+        return jnp.array([y, 1.0])
+
+    cond_fns = [cond_fn_1, cond_fn_2]
+    for cond_fn in cond_fns:
+        event = diffrax.Event(cond_fn=cond_fn)
+        with pytest.raises(ValueError):
+            diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
+
+
+def test_event_dtype_error():
+    term = diffrax.ODETerm(lambda t, y, args: y)
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = jnp.inf
+    dt0 = 1
+    y0 = 1.0
+
+    def cond_fn_1(t, y, args, **kwargs):
+        del t, args, kwargs
+        assert isinstance(y, jax.Array)
+        return jnp.array(1, dtype=int)
+
+    def cond_fn_2(t, y, args, **kwargs):
+        del t, args, kwargs
+        return jnp.array(1, dtype=complex)
+
+    cond_fns = [cond_fn_1, cond_fn_2]
+    for cond_fn in cond_fns:
+        event = diffrax.Event(cond_fn=cond_fn)
+        with pytest.raises(AssertionError):
+            diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)

From 57d90c520add1d832226c8fd5e99b74d0b83e345 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Mon, 27 May 2024 14:48:55 +0200
Subject: [PATCH 13/28] Fixed small error in the save_index update for events

---
 diffrax/_integrate.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index e29f00cf..a63f0e45 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -729,13 +729,14 @@ def _update_index(save_state):
 
         save_state = jtu.tree_map(
             lambda s: _update_index(s),
-            save_state,
+            final_state.save_state,
             is_leaf=lambda s: isinstance(s, SaveState),
         )
         final_state = eqx.tree_at(
             lambda s: s.save_state,
             final_state,
             save_state,
+            is_leaf=_is_none,
         )
 
     def _save_t1(subsaveat, save_state):

From 1bdf1d21c8379b4997747fb705ef1ab4b9c4fc07 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Mon, 27 May 2024 19:09:37 +0200
Subject: [PATCH 14/28] Updated how events are saved

When passing `SaveAt(steps=True, ts=ts)` for some array `ts` values will be saved at the times in `ts` in the time increments of each step of the solver. In practice this means that some of the saved values might be after the event time. I changed it so that these values are deleted.
---
 diffrax/_integrate.py | 75 ++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 67 insertions(+), 8 deletions(-)

diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index a63f0e45..6f3f23c3 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -721,17 +721,75 @@ def _call_real_impl():
             lambda: (final_state.tprev, final_state.y, result),
         )
 
-        # Update save_index to replace last saved step with event values
-        def _update_index(save_state):
-            return eqx.tree_at(
-                lambda s: s.save_index, save_state, replace_fn=lambda i: i - 1
-            )
+        # We delete all the saved values after the event time.
+        # For values saved at steps
+        def unsave_step(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState:
+            if subsaveat.steps:
+                save_index = save_state.save_index - 1
+                _ts = save_state.ts.at[save_index].set(jnp.inf)
+                _ys = jtu.tree_map(
+                    lambda _, __ys: __ys.at[save_index].set(jnp.inf),
+                    subsaveat.fn(tfinal, yfinal, args),
+                    save_state.ys,
+                )
+                save_state = eqx.tree_at(
+                    lambda s: s.save_index, save_state, replace_fn=lambda i: i - 1
+                )
+                save_state = SaveState(
+                    saveat_ts_index=save_state.saveat_ts_index,
+                    ts=_ts,
+                    ys=_ys,
+                    save_index=save_index,
+                )
+            return save_state
 
         save_state = jtu.tree_map(
-            lambda s: _update_index(s),
+            unsave_step,
+            saveat.subs,
             final_state.save_state,
-            is_leaf=lambda s: isinstance(s, SaveState),
+            is_leaf=_is_subsaveat,
         )
+
+        # For values saved at specific times, ts
+        def unsave_ts(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState:
+            if subsaveat.ts is not None:
+                save_state = unsave_ts_impl(subsaveat.ts, subsaveat.fn, save_state)
+            return save_state
+
+        def unsave_ts_impl(ts, fn, save_state: SaveState) -> SaveState:
+            def _cond_fun(_save_state):
+                return (ts[_save_state.saveat_ts_index - 1] > tfinal) & (
+                    _save_state.saveat_ts_index - 1 > 0
+                )
+
+            def _body_fun(_save_state):
+                saveat_ts_index = _save_state.saveat_ts_index - 1
+                _ts = _save_state.ts.at[saveat_ts_index].set(jnp.inf)
+                _ys = jtu.tree_map(
+                    lambda _, __ys: __ys.at[saveat_ts_index].set(jnp.inf),
+                    fn(tfinal, yfinal, args),
+                    _save_state.ys,
+                )
+                return SaveState(
+                    saveat_ts_index=saveat_ts_index,
+                    ts=_ts,
+                    ys=_ys,
+                    save_index=_save_state.save_index - 1,
+                )
+
+            return inner_while_loop(
+                _cond_fun,
+                _body_fun,
+                save_state,
+                max_steps=len(ts),
+                buffers=_inner_buffers,
+                checkpoints=len(ts),
+            )
+
+        save_state = jtu.tree_map(
+            unsave_ts, saveat.subs, save_state, is_leaf=_is_subsaveat
+        )
+
         final_state = eqx.tree_at(
             lambda s: s.save_state,
             final_state,
@@ -747,7 +805,8 @@ def _save_t1(subsaveat, save_state):
         else:
             if subsaveat.t1 or subsaveat.steps:
                 # In this branch we need to replace the last value with tfinal
-                # and yfinal returned by the root finder also if subsaveat.steps.
+                # and yfinal returned by the root finder also if subsaveat.steps
+                # because we deled the last value after the event time above.
                 save_state = _save(tfinal, yfinal, args, subsaveat.fn, save_state)
         return save_state
 

From 55e04e816f12e3595baf1a38625b53a2fb0e6879 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Mon, 27 May 2024 19:18:43 +0200
Subject: [PATCH 15/28] Added tests for different configurations of saveat

---
 test/test_event.py | 149 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 149 insertions(+)

diff --git a/test/test_event.py b/test/test_event.py
index 6214c575..aea42c41 100644
--- a/test/test_event.py
+++ b/test/test_event.py
@@ -2,6 +2,7 @@
 from typing import cast
 
 import diffrax
+import equinox as eqx
 import jax
 import jax.numpy as jnp
 import optimistix as optx
@@ -560,3 +561,151 @@ def cond_fn_2(t, y, args, **kwargs):
         event = diffrax.Event(cond_fn=cond_fn)
         with pytest.raises(AssertionError):
             diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
+
+
+@pytest.mark.parametrize("steps", (1, 2, 3, 4, 5))
+def test_event_save_steps(steps):
+    term = diffrax.ODETerm(lambda t, y, args: (1.0, 1.0))
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = 10
+    dt0 = 1
+    thr = steps - 0.5
+    y0 = (0.0, -thr)
+    ts = jnp.array([0.5, 3.5, 5.5])
+
+    def cond_fn(t, y, args, **kwargs):
+        del t, args, kwargs
+        x, _ = y
+        return x - thr
+
+    root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
+    event = diffrax.Event(cond_fn, root_finder)
+
+    def run(saveat):
+        sol = diffrax.diffeqsolve(
+            term,
+            solver,
+            t0,
+            t1,
+            dt0,
+            y0,
+            event=event,
+            saveat=saveat,
+        )
+        return cast(Array, sol.ts), cast(tuple, sol.ys)
+
+    saveats = [
+        diffrax.SaveAt(steps=True),
+        diffrax.SaveAt(steps=True, t1=True),
+        diffrax.SaveAt(steps=True, t1=True, t0=True),
+        diffrax.SaveAt(steps=True, fn=lambda t, y, args: (y[0], y[1] + thr)),
+    ]
+    num_steps = [steps, steps, steps + 1, steps]
+    yevents = [(thr, 0), (thr, 0), (thr, 0), (thr, thr)]
+
+    for saveat, n, yevent in zip(saveats, num_steps, yevents):
+        ts, ys = run(saveat)
+        xs, zs = ys
+        xevent, zevent = yevent
+        assert jnp.sum(jnp.isfinite(ts)) == n
+        assert jnp.sum(jnp.isfinite(xs)) == n
+        assert jnp.sum(jnp.isfinite(zs)) == n
+        assert jnp.all(jnp.isclose(ts[n - 1], thr, atol=1e-5))
+        assert jnp.all(jnp.isclose(xs[n - 1], xevent, atol=1e-5))
+        assert jnp.all(jnp.isclose(zs[n - 1], zevent, atol=1e-5))
+
+
+@pytest.mark.parametrize("steps", (1, 2, 3, 4, 5))
+def test_event_save_ts(steps):
+    term = diffrax.ODETerm(lambda t, y, args: (1.0, 1.0))
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = 10
+    dt0 = 1
+    thr = steps - 0.5
+    y0 = (0.0, -thr)
+    ts = jnp.array([0.5, 3.5, 5.5])
+
+    def cond_fn(t, y, args, **kwargs):
+        del t, args, kwargs
+        x, _ = y
+        return x - thr
+
+    root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
+    event = diffrax.Event(cond_fn, root_finder)
+
+    def run(saveat):
+        sol = diffrax.diffeqsolve(
+            term,
+            solver,
+            t0,
+            t1,
+            dt0,
+            y0,
+            event=event,
+            saveat=saveat,
+        )
+        return cast(Array, sol.ts), cast(tuple, sol.ys)
+
+    saveats = [
+        diffrax.SaveAt(ts=ts),
+        diffrax.SaveAt(ts=ts, t1=True),
+        diffrax.SaveAt(ts=ts, t0=True),
+        diffrax.SaveAt(ts=ts, steps=True),
+        diffrax.SaveAt(ts=ts, fn=lambda t, y, args: (y[0], y[1] + thr)),
+    ]
+    save_finals = [False, True, False, True, False]
+    yevents = [(thr, 0), (thr, 0), (thr, 0), (thr, 0), (thr, thr)]
+    for saveat, save_final, yevent in zip(saveats, save_finals, yevents):
+        ts, ys = run(saveat)
+        xs, zs = ys
+        xevent, zevent = yevent
+        if save_final:
+            assert jnp.all(jnp.isclose(ts[jnp.isfinite(ts)][-1], thr, atol=1e-5))
+            assert jnp.all(jnp.isclose(xs[jnp.isfinite(xs)][-1], xevent, atol=1e-5))
+            assert jnp.all(jnp.isclose(zs[jnp.isfinite(zs)][-1], zevent, atol=1e-5))
+        else:
+            assert jnp.all(ts[jnp.isfinite(ts)] <= thr)
+
+
+@pytest.mark.parametrize("steps", (1, 2, 3, 4, 5))
+def test_event_save_subsaveat(steps):
+    term = diffrax.ODETerm(lambda t, y, args: jnp.array([1.0, 1.0]))
+    solver = diffrax.Tsit5()
+    t0 = 0.0
+    t1 = 10.0
+    dt0 = 1.0
+    thr = steps - 0.5
+    y0 = jnp.array([0.0, -thr])
+    ts = jnp.arange(t0, t1, 3.0)
+    ts_event = jnp.sum(ts <= thr)
+    last_t = jnp.array(ts[ts_event - 1])
+
+    def cond_fn(t, y, args, **kwargs):
+        del t, args, kwargs
+        x, _ = y
+        return x - thr
+
+    root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
+    event = diffrax.Event(cond_fn, root_finder)
+
+    class Saved(eqx.Module):
+        y: Array
+
+    def save_fn(t, y, args):
+        del t, args
+        ynorm = jnp.einsum("i,i->", y, y)
+        return Saved(jnp.array([ynorm]))
+
+    last_save = save_fn(None, y0 + last_t, None).y
+    subsaveat_a = diffrax.SubSaveAt(ts=ts, fn=save_fn)
+    subsaveat_b = diffrax.SubSaveAt(steps=True)
+    saveat = diffrax.SaveAt(subs=[subsaveat_a, subsaveat_b])
+    sol = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event, saveat=saveat)
+    ts_1, ts_2 = cast(list, sol.ts)
+    ys_1, ys_2 = cast(list, sol.ys)
+    assert jnp.sum(jnp.isfinite(ts_1)) == ts_event
+    assert jnp.sum(jnp.isfinite(ts_2)) == steps
+    assert jnp.all(jnp.isclose(ys_2[steps - 1], jnp.array([thr, 0]), atol=1e-5))
+    assert jnp.all(jnp.isclose(ys_1.y[ts_event - 1], last_save, atol=1e-5))

From d8a8ba7734796a9784502ad86a92dd34127066d4 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Tue, 28 May 2024 10:07:33 +0200
Subject: [PATCH 16/28] Changed to ValueError when cond_fn returns
 non-boolean/float.

---
 diffrax/_integrate.py | 10 ++++++++--
 test/test_event.py    |  2 +-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index 6f3f23c3..5930222f 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -538,7 +538,10 @@ def _outer_cond_fn(cond_fn_i, old_event_value_i):
                 elif jnp.issubdtype(new_dtype, jnp.bool_):
                     event_mask_i = new_event_value_i
                 else:
-                    assert False
+                    raise ValueError(
+                        "Event functions must return either a boolean or a float, got "
+                        f"{new_dtype}."
+                    )
                 return new_event_value_i, event_mask_i
 
             event_values__mask = jtu.tree_map(
@@ -1317,7 +1320,10 @@ def _outer_cond_fn(cond_fn_i):
             elif jnp.issubdtype(event_dtype, jnp.bool_):
                 event_mask_i = event_value_i
             else:
-                assert False
+                raise ValueError(
+                    "Event functions must return either a boolean or a float, got "
+                    f"{event_dtype}."
+                )
             return event_value_i, event_mask_i
 
         event_values__mask = jtu.tree_map(
diff --git a/test/test_event.py b/test/test_event.py
index aea42c41..dbfcd255 100644
--- a/test/test_event.py
+++ b/test/test_event.py
@@ -559,7 +559,7 @@ def cond_fn_2(t, y, args, **kwargs):
     cond_fns = [cond_fn_1, cond_fn_2]
     for cond_fn in cond_fns:
         event = diffrax.Event(cond_fn=cond_fn)
-        with pytest.raises(AssertionError):
+        with pytest.raises(ValueError):
             diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
 
 

From 70e044febe20d3b3de19cf031776376ef88b8b83 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Tue, 28 May 2024 10:28:13 +0200
Subject: [PATCH 17/28] Added docstring to Event class

---
 diffrax/_event.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/diffrax/_event.py b/diffrax/_event.py
index 54c1ef9f..e2ef42f2 100644
--- a/diffrax/_event.py
+++ b/diffrax/_event.py
@@ -11,10 +11,31 @@
 
 
 class Event(eqx.Module):
+    """Can be used to terminate the solve early if one of multiple conditions
+    is triggered. It allows for both continuous and boolean condition functions. In the
+    former case, a root finder can be used to find the exact time of the event.
+
+    Instances of this class should be passed as the `event` argument of
+    [`diffrax.diffeqsolve`][].
+    """
+
     cond_fn: PyTree[Callable[..., Union[BoolScalarLike, RealScalarLike]]]
     root_finder: Optional[optx.AbstractRootFinder] = None
 
 
+Event.__init__.__doc__ = """**Arguments:**
+
+- `cond_fn`: A PyTree of functions `f(t, y, args, **kwargs) -> c` returning a boolean or
+    a real number. If the return value is a boolean, the solve will terminate when `c`
+    is `True`. If the return value is a real number, the solve will terminate when `c`
+    changes sign.
+- `root_finder`: An optional root finder to use for finding the exact time of the event.
+    If the triggered condition function is boolean,  the returned time will be the right
+    endpoint of the last successful step.
+
+"""
+
+
 def steady_state_event(
     rtol: Optional[float] = None,
     atol: Optional[float] = None,

From 4c509b6c22ecbdd562d2633fdbfa392438595bb8 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Tue, 28 May 2024 10:32:57 +0200
Subject: [PATCH 18/28] Updated docstring for steady_state_event

---
 diffrax/_event.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/diffrax/_event.py b/diffrax/_event.py
index e2ef42f2..d5026ece 100644
--- a/diffrax/_event.py
+++ b/diffrax/_event.py
@@ -41,8 +41,9 @@ def steady_state_event(
     atol: Optional[float] = None,
     norm: Optional[Callable[[PyTree[Array]], RealScalarLike]] = None,
 ):
-    """Create a [`diffrax.Event`][] that terminates the solve once a steady state is
-    achieved.
+    """Create a condition function that terminates the solve once a steady state is
+    achieved. The returned function should be passed as the `cond_fn` argument of
+    [`diffrax.Event`][].
 
     **Arguments:**
 
@@ -53,8 +54,8 @@ def steady_state_event(
 
     **Returns:**
 
-    A [`diffrax.Event`][] object, that can be passed to
-    `diffrax.diffeqsolve(..., event=...)`.
+    A function `f(t, y, args, **kwargs)`, that can be passed to
+    `diffrax.Evetnt(cond_fn=..., ...)`.
     """
 
     def _cond_fn(t, y, args, *, terms, solver, stepsize_controller, **kwargs):

From fbea7943dd8aa76b2cdca0d000819cff01e9680d Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Tue, 28 May 2024 10:45:31 +0200
Subject: [PATCH 19/28] Updated docstring for ImplicitAdjoint

---
 diffrax/_adjoint.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/diffrax/_adjoint.py b/diffrax/_adjoint.py
index c75d0a9e..4ff2dd2c 100644
--- a/diffrax/_adjoint.py
+++ b/diffrax/_adjoint.py
@@ -450,7 +450,8 @@ class ImplicitAdjoint(AbstractAdjoint):
     r"""Backpropagate via the [implicit function theorem](https://en.wikipedia.org/wiki/Implicit_function_theorem#Statement_of_the_theorem).
 
     This is used when solving towards a steady state, typically using
-    [`diffrax.SteadyStateEvent`][]. In this case, the output of the solver is $y(θ)$
+    [`diffrax.Event`][] where the condition function is obtained by calling
+    [`diffrax.steady_state_event`][]. In this case, the output of the solver is $y(θ)$
     for which $f(t, y(θ), θ) = 0$. (Where $θ$ corresponds to all parameters found
     through `terms` and `args`, but not `y0`.) Then we can skip backpropagating through
     the solver and instead directly compute

From b158800235894c4bef5c9de545aab3b4ea966c81 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Tue, 28 May 2024 11:36:19 +0200
Subject: [PATCH 20/28] Added example to Event docstring

---
 diffrax/_event.py | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/diffrax/_event.py b/diffrax/_event.py
index d5026ece..72b2aa3f 100644
--- a/diffrax/_event.py
+++ b/diffrax/_event.py
@@ -33,6 +33,46 @@ class Event(eqx.Module):
     If the triggered condition function is boolean,  the returned time will be the right
     endpoint of the last successful step.
 
+!!! Example
+
+    Consider a bouncing ball dropped from some intial height $x_0$. We can model 
+    the ball by a 2-dimensional ODE
+
+    $dx_t = v_t dt, \\quad dv_t = -g dt,$
+
+    where $x_t$ represents the height of the ball, $v_t$ its velocity, 
+    and $g$ is the gravitational constant. With $g=8$, this corresponds to the
+    vector field:
+    ```python
+    def vf(t, y, args):
+        _, v = y
+        return jnp.array([v, -8.0])
+    ```
+
+    Figuring out exactly when the ball hits the ground amounts to 
+    solving the ODE until the event $x_t=0$ is triggered. This can be done by using 
+    the real-valued condition function:
+    ```python
+    def cond_fn(t, y, args, **kwargs):
+        x, _ = y
+        return x
+    ```
+
+    With $x_0=10$, this would yield:
+    ```python
+    y0 = jnp.array([10.0, 0.0])
+    t0 = 0
+    t1 = jnp.inf
+    dt0 = 0.1
+    term = diffrax.ODETerm(vector_field)
+    root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
+    event = diffrax.Event(cond_fn, root_finder)
+    solver = diffrax.Tsit5()
+    sol = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event, saveat=saveat)
+    print(f"Event time: {sol.ts[0]}") # Event time: 1.58...
+    print(f"Velocity at event time: {sol.ys[0, 1]}") # Velocity at event time: -12.64...
+    ```
+
 """
 
 

From 812e5c667110e48c90d41fd2be2b6f3b17092469 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Tue, 28 May 2024 11:44:08 +0200
Subject: [PATCH 21/28] Updated steady state example to use the new syntax

---
 examples/steady_state.ipynb | 209 ++++++++++++++++++------------------
 1 file changed, 105 insertions(+), 104 deletions(-)

diff --git a/examples/steady_state.ipynb b/examples/steady_state.ipynb
index 1cbd1653..dac8aaa3 100644
--- a/examples/steady_state.ipynb
+++ b/examples/steady_state.ipynb
@@ -102,7 +102,8 @@
     "    y0 = 1.0\n",
     "    max_steps = None\n",
     "    controller = diffrax.PIDController(rtol=1e-3, atol=1e-6)\n",
-    "    event = diffrax.SteadyStateEvent()\n",
+    "    cond_fn = diffrax.steady_state_event()\n",
+    "    event = diffrax.Event(cond_fn)\n",
     "    adjoint = diffrax.ImplicitAdjoint()\n",
     "    # This combination of event, t1, max_steps, adjoint is particularly\n",
     "    # natural: we keep integration forever until we hit the event, with\n",
@@ -117,7 +118,7 @@
     "        y0,\n",
     "        max_steps=max_steps,\n",
     "        stepsize_controller=controller,\n",
-    "        discrete_terminating_event=event,\n",
+    "        event=event,\n",
     "        adjoint=adjoint,\n",
     "    )\n",
     "    (y1,) = sol.ys\n",
@@ -148,105 +149,105 @@
      "output_type": "stream",
      "text": [
       "Step: 0 Steady State: 0.025839969515800476\n",
-      "Step: 1 Steady State: 0.058249037712812424\n",
-      "Step: 2 Steady State: 0.09451574087142944\n",
-      "Step: 3 Steady State: 0.13270404934883118\n",
-      "Step: 4 Steady State: 0.17144456505775452\n",
-      "Step: 5 Steady State: 0.2097906768321991\n",
-      "Step: 6 Steady State: 0.24709917604923248\n",
-      "Step: 7 Steady State: 0.28294336795806885\n",
-      "Step: 8 Steady State: 0.3170691728591919\n",
-      "Step: 9 Steady State: 0.34933507442474365\n",
-      "Step: 10 Steady State: 0.37968066334724426\n",
-      "Step: 11 Steady State: 0.4081019163131714\n",
-      "Step: 12 Steady State: 0.43463483452796936\n",
-      "Step: 13 Steady State: 0.45934173464775085\n",
-      "Step: 14 Steady State: 0.4823019802570343\n",
-      "Step: 15 Steady State: 0.5035936236381531\n",
-      "Step: 16 Steady State: 0.5233209133148193\n",
-      "Step: 17 Steady State: 0.5415788888931274\n",
-      "Step: 18 Steady State: 0.5584676265716553\n",
-      "Step: 19 Steady State: 0.5740787982940674\n",
-      "Step: 20 Steady State: 0.5885017514228821\n",
-      "Step: 21 Steady State: 0.6018210053443909\n",
-      "Step: 22 Steady State: 0.6141175627708435\n",
-      "Step: 23 Steady State: 0.6254667043685913\n",
-      "Step: 24 Steady State: 0.6359376907348633\n",
-      "Step: 25 Steady State: 0.6455990076065063\n",
-      "Step: 26 Steady State: 0.6545112729072571\n",
-      "Step: 27 Steady State: 0.6627309322357178\n",
-      "Step: 28 Steady State: 0.6703115701675415\n",
-      "Step: 29 Steady State: 0.6773026585578918\n",
-      "Step: 30 Steady State: 0.6837494373321533\n",
-      "Step: 31 Steady State: 0.6896938681602478\n",
-      "Step: 32 Steady State: 0.6951748728752136\n",
-      "Step: 33 Steady State: 0.7002284526824951\n",
-      "Step: 34 Steady State: 0.7048872113227844\n",
-      "Step: 35 Steady State: 0.7091819047927856\n",
-      "Step: 36 Steady State: 0.7131412029266357\n",
-      "Step: 37 Steady State: 0.7167739868164062\n",
-      "Step: 38 Steady State: 0.7201183438301086\n",
-      "Step: 39 Steady State: 0.7231980562210083\n",
-      "Step: 40 Steady State: 0.7260348796844482\n",
-      "Step: 41 Steady State: 0.7286462187767029\n",
-      "Step: 42 Steady State: 0.7310511469841003\n",
-      "Step: 43 Steady State: 0.733269989490509\n",
-      "Step: 44 Steady State: 0.7353137731552124\n",
-      "Step: 45 Steady State: 0.7371994853019714\n",
-      "Step: 46 Steady State: 0.7389383912086487\n",
-      "Step: 47 Steady State: 0.740541934967041\n",
-      "Step: 48 Steady State: 0.7420334219932556\n",
-      "Step: 49 Steady State: 0.7434003353118896\n",
-      "Step: 50 Steady State: 0.7446598410606384\n",
-      "Step: 51 Steady State: 0.7458205819129944\n",
-      "Step: 52 Steady State: 0.7468900680541992\n",
-      "Step: 53 Steady State: 0.7478761672973633\n",
-      "Step: 54 Steady State: 0.7487852573394775\n",
-      "Step: 55 Steady State: 0.7496234178543091\n",
-      "Step: 56 Steady State: 0.750394344329834\n",
-      "Step: 57 Steady State: 0.7511063814163208\n",
-      "Step: 58 Steady State: 0.751763105392456\n",
-      "Step: 59 Steady State: 0.7523672580718994\n",
-      "Step: 60 Steady State: 0.7529228329658508\n",
-      "Step: 61 Steady State: 0.753433346748352\n",
-      "Step: 62 Steady State: 0.7539049983024597\n",
-      "Step: 63 Steady State: 0.7543382048606873\n",
-      "Step: 64 Steady State: 0.7547407746315002\n",
-      "Step: 65 Steady State: 0.7551127672195435\n",
-      "Step: 66 Steady State: 0.7554563879966736\n",
-      "Step: 67 Steady State: 0.7557693123817444\n",
-      "Step: 68 Steady State: 0.7560611367225647\n",
-      "Step: 69 Steady State: 0.7563308477401733\n",
-      "Step: 70 Steady State: 0.7565800547599792\n",
-      "Step: 71 Steady State: 0.756810188293457\n",
-      "Step: 72 Steady State: 0.7570226788520813\n",
-      "Step: 73 Steady State: 0.7572163343429565\n",
-      "Step: 74 Steady State: 0.7573966979980469\n",
-      "Step: 75 Steady State: 0.7575633525848389\n",
-      "Step: 76 Steady State: 0.7577127814292908\n",
-      "Step: 77 Steady State: 0.7578537464141846\n",
-      "Step: 78 Steady State: 0.7579842805862427\n",
-      "Step: 79 Steady State: 0.7581048607826233\n",
-      "Step: 80 Steady State: 0.7582123279571533\n",
-      "Step: 81 Steady State: 0.7583134770393372\n",
-      "Step: 82 Steady State: 0.7584078907966614\n",
-      "Step: 83 Steady State: 0.7584953904151917\n",
-      "Step: 84 Steady State: 0.758575975894928\n",
-      "Step: 85 Steady State: 0.7586501836776733\n",
-      "Step: 86 Steady State: 0.7587193250656128\n",
-      "Step: 87 Steady State: 0.7587832808494568\n",
-      "Step: 88 Steady State: 0.7588424682617188\n",
-      "Step: 89 Steady State: 0.7588958144187927\n",
-      "Step: 90 Steady State: 0.7589460015296936\n",
-      "Step: 91 Steady State: 0.7589924931526184\n",
-      "Step: 92 Steady State: 0.7590354681015015\n",
-      "Step: 93 Steady State: 0.7590752243995667\n",
-      "Step: 94 Steady State: 0.7591111063957214\n",
-      "Step: 95 Steady State: 0.7591448426246643\n",
-      "Step: 96 Steady State: 0.7591760754585266\n",
-      "Step: 97 Steady State: 0.7592049241065979\n",
-      "Step: 98 Steady State: 0.7592315673828125\n",
-      "Step: 99 Steady State: 0.7592562437057495\n",
+      "Step: 1 Steady State: 0.05824900045990944\n",
+      "Step: 2 Steady State: 0.09451568126678467\n",
+      "Step: 3 Steady State: 0.1327039748430252\n",
+      "Step: 4 Steady State: 0.1714443564414978\n",
+      "Step: 5 Steady State: 0.20979028940200806\n",
+      "Step: 6 Steady State: 0.24709881842136383\n",
+      "Step: 7 Steady State: 0.28294941782951355\n",
+      "Step: 8 Steady State: 0.31707584857940674\n",
+      "Step: 9 Steady State: 0.34934186935424805\n",
+      "Step: 10 Steady State: 0.37968698143959045\n",
+      "Step: 11 Steady State: 0.4081074893474579\n",
+      "Step: 12 Steady State: 0.43463948369026184\n",
+      "Step: 13 Steady State: 0.45934492349624634\n",
+      "Step: 14 Steady State: 0.48230400681495667\n",
+      "Step: 15 Steady State: 0.5036059021949768\n",
+      "Step: 16 Steady State: 0.5233321189880371\n",
+      "Step: 17 Steady State: 0.5415896773338318\n",
+      "Step: 18 Steady State: 0.5584752559661865\n",
+      "Step: 19 Steady State: 0.5740804076194763\n",
+      "Step: 20 Steady State: 0.5884985327720642\n",
+      "Step: 21 Steady State: 0.6018134951591492\n",
+      "Step: 22 Steady State: 0.6141058206558228\n",
+      "Step: 23 Steady State: 0.6254505515098572\n",
+      "Step: 24 Steady State: 0.6359192728996277\n",
+      "Step: 25 Steady State: 0.6455777287483215\n",
+      "Step: 26 Steady State: 0.6544871926307678\n",
+      "Step: 27 Steady State: 0.6627050638198853\n",
+      "Step: 28 Steady State: 0.6702842116355896\n",
+      "Step: 29 Steady State: 0.6772737503051758\n",
+      "Step: 30 Steady State: 0.6837191581726074\n",
+      "Step: 31 Steady State: 0.6896624565124512\n",
+      "Step: 32 Steady State: 0.6951420903205872\n",
+      "Step: 33 Steady State: 0.7001940608024597\n",
+      "Step: 34 Steady State: 0.7048525214195251\n",
+      "Step: 35 Steady State: 0.709147572517395\n",
+      "Step: 36 Steady State: 0.7131075263023376\n",
+      "Step: 37 Steady State: 0.7167584300041199\n",
+      "Step: 38 Steady State: 0.720124363899231\n",
+      "Step: 39 Steady State: 0.7232275605201721\n",
+      "Step: 40 Steady State: 0.7260884642601013\n",
+      "Step: 41 Steady State: 0.7287259697914124\n",
+      "Step: 42 Steady State: 0.7311574816703796\n",
+      "Step: 43 Steady State: 0.7333983778953552\n",
+      "Step: 44 Steady State: 0.7354647517204285\n",
+      "Step: 45 Steady State: 0.7373697757720947\n",
+      "Step: 46 Steady State: 0.7391260266304016\n",
+      "Step: 47 Steady State: 0.7407451272010803\n",
+      "Step: 48 Steady State: 0.7422377467155457\n",
+      "Step: 49 Steady State: 0.7436137795448303\n",
+      "Step: 50 Steady State: 0.7448822855949402\n",
+      "Step: 51 Steady State: 0.7460517287254333\n",
+      "Step: 52 Steady State: 0.7471297979354858\n",
+      "Step: 53 Steady State: 0.7481234669685364\n",
+      "Step: 54 Steady State: 0.7490396499633789\n",
+      "Step: 55 Steady State: 0.7498842477798462\n",
+      "Step: 56 Steady State: 0.7506628632545471\n",
+      "Step: 57 Steady State: 0.7513806223869324\n",
+      "Step: 58 Steady State: 0.7520219683647156\n",
+      "Step: 59 Steady State: 0.7526065707206726\n",
+      "Step: 60 Steady State: 0.7531405687332153\n",
+      "Step: 61 Steady State: 0.7536292672157288\n",
+      "Step: 62 Steady State: 0.754077136516571\n",
+      "Step: 63 Steady State: 0.7544881105422974\n",
+      "Step: 64 Steady State: 0.7548655867576599\n",
+      "Step: 65 Steady State: 0.7552322149276733\n",
+      "Step: 66 Steady State: 0.7555564045906067\n",
+      "Step: 67 Steady State: 0.7558530569076538\n",
+      "Step: 68 Steady State: 0.7561249732971191\n",
+      "Step: 69 Steady State: 0.7563938498497009\n",
+      "Step: 70 Steady State: 0.7566279768943787\n",
+      "Step: 71 Steady State: 0.7568415403366089\n",
+      "Step: 72 Steady State: 0.7570368051528931\n",
+      "Step: 73 Steady State: 0.7572155594825745\n",
+      "Step: 74 Steady State: 0.7573794722557068\n",
+      "Step: 75 Steady State: 0.7575299143791199\n",
+      "Step: 76 Steady State: 0.757668137550354\n",
+      "Step: 77 Steady State: 0.7577952742576599\n",
+      "Step: 78 Steady State: 0.7579122185707092\n",
+      "Step: 79 Steady State: 0.7580198645591736\n",
+      "Step: 80 Steady State: 0.7581189870834351\n",
+      "Step: 81 Steady State: 0.758210301399231\n",
+      "Step: 82 Steady State: 0.7583132982254028\n",
+      "Step: 83 Steady State: 0.7583956122398376\n",
+      "Step: 84 Steady State: 0.7584698796272278\n",
+      "Step: 85 Steady State: 0.7585371136665344\n",
+      "Step: 86 Steady State: 0.7585982084274292\n",
+      "Step: 87 Steady State: 0.7586538791656494\n",
+      "Step: 88 Steady State: 0.7587047219276428\n",
+      "Step: 89 Steady State: 0.7587512731552124\n",
+      "Step: 90 Steady State: 0.7587938904762268\n",
+      "Step: 91 Steady State: 0.7588329911231995\n",
+      "Step: 92 Steady State: 0.758868932723999\n",
+      "Step: 93 Steady State: 0.7589019536972046\n",
+      "Step: 94 Steady State: 0.7589322924613953\n",
+      "Step: 95 Steady State: 0.7589602470397949\n",
+      "Step: 96 Steady State: 0.7589859366416931\n",
+      "Step: 97 Steady State: 0.7590096592903137\n",
+      "Step: 98 Steady State: 0.7590314745903015\n",
+      "Step: 99 Steady State: 0.7590516209602356\n",
       "Target: 0.7599999904632568\n"
      ]
     }
@@ -278,9 +279,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "py37",
+   "display_name": "diffrax",
    "language": "python",
-   "name": "py37"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -292,7 +293,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.13"
+   "version": "3.11.0"
   }
  },
  "nbformat": 4,

From eb16e57c8e521f5915d6ee98515c5b3d1680a804 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Sat, 1 Jun 2024 20:41:05 +0200
Subject: [PATCH 22/28] Fixed weird type checker error

---
 diffrax/_integrate.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index 5930222f..10d71d86 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -553,7 +553,7 @@ def _outer_cond_fn(cond_fn_i, old_event_value_i):
             event_structure = jtu.tree_structure(event.cond_fn, is_leaf=callable)
             event_values, event_mask = jtu.tree_transpose(
                 event_structure,
-                None,  # pyright: ignore
+                jtu.tree_structure((0, 0)),
                 event_values__mask,
             )
             had_event = False
@@ -1334,7 +1334,7 @@ def _outer_cond_fn(cond_fn_i):
         event_structure = jtu.tree_structure(event.cond_fn, is_leaf=callable)
         event_values, event_mask = jtu.tree_transpose(
             event_structure,
-            None,  # pyright: ignore
+            jtu.tree_structure((0, 0)),
             event_values__mask,
         )
         had_event = False

From 09c92c3e739c208443adc870fb64a6f47a909a96 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Mon, 10 Jun 2024 09:36:12 +0200
Subject: [PATCH 23/28] Updated steady state test to use the new syntax

---
 test/test_adjoint.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/test/test_adjoint.py b/test/test_adjoint.py
index 834b7a3b..a6148adb 100644
--- a/test/test_adjoint.py
+++ b/test/test_adjoint.py
@@ -276,7 +276,8 @@ def loss(model, target_steady_state):
         y0 = 1.0
         max_steps = None
         controller = diffrax.PIDController(rtol=1e-3, atol=1e-6)
-        event = diffrax.SteadyStateEvent()
+        cond_fn = diffrax.steady_state_event()
+        event = diffrax.Event(cond_fn)
         adjoint = diffrax.ImplicitAdjoint()
         sol = diffrax.diffeqsolve(
             term,
@@ -287,7 +288,7 @@ def loss(model, target_steady_state):
             y0,
             max_steps=max_steps,
             stepsize_controller=controller,
-            discrete_terminating_event=event,
+            event=event,
             adjoint=adjoint,
         )
         (y1,) = cast(Array, sol.ys)

From d07c8f4c3513c130accf336db58b0e6c96e2a1cb Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Sat, 15 Jun 2024 12:43:09 +0200
Subject: [PATCH 24/28] Doc tweaks for events

---
 diffrax/_event.py     | 37 +++++++++++++++++++++++--------------
 diffrax/_integrate.py |  6 +++---
 diffrax/_solution.py  | 16 ++++++++--------
 docs/api/solution.md  |  6 ++++++
 4 files changed, 40 insertions(+), 25 deletions(-)

diff --git a/diffrax/_event.py b/diffrax/_event.py
index 72b2aa3f..3168b297 100644
--- a/diffrax/_event.py
+++ b/diffrax/_event.py
@@ -11,9 +11,10 @@
 
 
 class Event(eqx.Module):
-    """Can be used to terminate the solve early if one of multiple conditions
-    is triggered. It allows for both continuous and boolean condition functions. In the
-    former case, a root finder can be used to find the exact time of the event.
+    """Can be used to terminate the solve early if a condition, or one of multiple
+    conditions, is triggered. It allows for both boolean and continuous condition
+    functions. In the latter case, a root finder can be used to find the exact time of
+    the event. Boolean and continuous conditions can be used together.
 
     Instances of this class should be passed as the `event` argument of
     [`diffrax.diffeqsolve`][].
@@ -25,26 +26,33 @@ class Event(eqx.Module):
 
 Event.__init__.__doc__ = """**Arguments:**
 
-- `cond_fn`: A PyTree of functions `f(t, y, args, **kwargs) -> c` returning a boolean or
-    a real number. If the return value is a boolean, the solve will terminate when `c`
-    is `True`. If the return value is a real number, the solve will terminate when `c`
+- `cond_fn`: A function or PyTree of functions `f(t, y, args, **kwargs) -> c` each
+    returning either a boolean or a real number. If the return value is a boolean, then
+    the solve will terminate on the first step on which `c` becomes `True`. If the
+    return value is a real number, then the solve will terminate on the step when `c`
     changes sign.
-- `root_finder`: An optional root finder to use for finding the exact time of the event.
-    If the triggered condition function is boolean,  the returned time will be the right
-    endpoint of the last successful step.
+
+- `root_finder`: An optional [root finder](../nonlinear_solver/) to use for finding
+    the exact time of the event. If the triggered condition function returns a real
+    number, then the final time will be the time at which that real number equals zero.
+    (If the triggered condition function returns a boolean, then the returned time will
+    just be the end of the step on which it becomes `True`.) 
+    [`optimistix.Newton`](https://docs.kidger.site/optimistix/api/root_find/#optimistix.Newton)
+    would be a typical choice here.
 
 !!! Example
 
     Consider a bouncing ball dropped from some intial height $x_0$. We can model 
     the ball by a 2-dimensional ODE
 
-    $dx_t = v_t dt, \\quad dv_t = -g dt,$
+    $\\frac{dx_t}{dt} = v_t, \\quad \\frac{dv_t}{dt} = -g,$
 
     where $x_t$ represents the height of the ball, $v_t$ its velocity, 
     and $g$ is the gravitational constant. With $g=8$, this corresponds to the
     vector field:
+
     ```python
-    def vf(t, y, args):
+    def vector_field(t, y, args):
         _, v = y
         return jnp.array([v, -8.0])
     ```
@@ -52,6 +60,7 @@ def vf(t, y, args):
     Figuring out exactly when the ball hits the ground amounts to 
     solving the ODE until the event $x_t=0$ is triggered. This can be done by using 
     the real-valued condition function:
+
     ```python
     def cond_fn(t, y, args, **kwargs):
         x, _ = y
@@ -59,6 +68,7 @@ def cond_fn(t, y, args, **kwargs):
     ```
 
     With $x_0=10$, this would yield:
+
     ```python
     y0 = jnp.array([10.0, 0.0])
     t0 = 0
@@ -68,11 +78,10 @@ def cond_fn(t, y, args, **kwargs):
     root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
     event = diffrax.Event(cond_fn, root_finder)
     solver = diffrax.Tsit5()
-    sol = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event, saveat=saveat)
+    sol = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
     print(f"Event time: {sol.ts[0]}") # Event time: 1.58...
     print(f"Velocity at event time: {sol.ys[0, 1]}") # Velocity at event time: -12.64...
     ```
-
 """
 
 
@@ -95,7 +104,7 @@ def steady_state_event(
     **Returns:**
 
     A function `f(t, y, args, **kwargs)`, that can be passed to
-    `diffrax.Evetnt(cond_fn=..., ...)`.
+    `diffrax.Event(cond_fn=..., ...)`.
     """
 
     def _cond_fn(t, y, args, *, terms, solver, stepsize_controller, **kwargs):
diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index 10d71d86..92fccf02 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -564,7 +564,7 @@ def _outer_cond_fn(cond_fn_i, old_event_value_i):
             event_mask = jtu.tree_unflatten(event_structure, event_mask_leaves)
             result = RESULTS.where(
                 had_event,
-                RESULTS.terminating_event_occurred,
+                RESULTS.event_occurred,
                 result,
             )
 
@@ -680,7 +680,7 @@ def _call_real_impl():
                     # to avoid this, I think.
                     _value = lax.cond(_event_mask_i, _call_real_impl, lambda: 0.0)
 
-                    # Third: if no events triggered at all, then has the root occur at
+                    # Third: if no events triggered at all, then have the root occur at
                     # the end of the last step (which will be the `t1` of the overall
                     # solve).
                     _value = jnp.where(event_happened, _value, _distance_from_t_end)
@@ -1345,7 +1345,7 @@ def _outer_cond_fn(cond_fn_i):
         event_mask = jtu.tree_unflatten(event_structure, event_mask_leaves)
         result = RESULTS.where(
             had_event,
-            RESULTS.terminating_event_occurred,
+            RESULTS.event_occurred,
             result,
         )
         del had_event, event_structure, event_mask_leaves, event_values__mask
diff --git a/diffrax/_solution.py b/diffrax/_solution.py
index f79b60c5..3c895d62 100644
--- a/diffrax/_solution.py
+++ b/diffrax/_solution.py
@@ -18,7 +18,7 @@ class RESULTS(optx.RESULTS):  # pyright: ignore
     dt_min_reached = (
         "The minimum step size was reached in the differential equation solver."
     )
-    terminating_event_occurred = (
+    event_occurred = (
         "Terminating differential equation solve because an event occurred."
     )
 
@@ -32,9 +32,10 @@ def discrete_terminating_event_occurred(self):
     warnings.warn(
         "`diffrax.RESULTS.discrete_terminating_event_occurred` is deprecated in "
         "favour of `diffrax.RESULTS.terminating_event_occurred`. This will be "
-        "removed in some future version of Diffrax."
+        "removed in some future version of Diffrax.",
+        stacklevel=2,
     )
-    return self.terminating_event_occurred
+    return self.event_occurred
 
 
 RESULTS.discrete_terminating_event_occurred = discrete_terminating_event_occurred  # pyright: ignore[reportAttributeAccessIssue]
@@ -49,10 +50,8 @@ def is_successful(result: RESULTS) -> Bool[Array, ""]:
     return result == RESULTS.successful
 
 
-# TODO: In the future we may support other event types, in which case this function
-# should be updated.
 def is_event(result: RESULTS) -> Bool[Array, ""]:
-    return result == RESULTS.terminating_event_occurred
+    return result == RESULTS.event_occurred
 
 
 def update_result(old_result: RESULTS, new_result: RESULTS) -> RESULTS:
@@ -84,8 +83,9 @@ class Solution(AbstractPath):
     - `ys`: The value of the solution at each of the times in `ts`. Might `None` if no
         values were saved.
     - `stats`: Statistics for the solve (number of steps etc.).
-    - `result`: Enumeration specifying the success or cause of failure of the solve.
-        A human-readable message is displayed if printed. No message means success!
+    - `result`: A [`diffrax.RESULT`][] specifying the success or cause of failure of the
+        solve. A human-readable message is displayed if printed. No message means
+        success!
     - `solver_state`: If saved, the final internal state of the numerical solver.
     - `controller_state`: If saved, the final internal state for the step size
         controller.
diff --git a/docs/api/solution.md b/docs/api/solution.md
index 5a74c90d..8682d8b9 100644
--- a/docs/api/solution.md
+++ b/docs/api/solution.md
@@ -15,3 +15,9 @@
             - message
             - evaluate
             - derivative
+
+---
+
+::: diffrax.RESULTS
+    selection:
+        members: false

From 883841fc7e807319b8a945237b7e33dcaf96a4cc Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Sun, 16 Jun 2024 13:17:40 +0200
Subject: [PATCH 25/28] Typo in comment

---
 diffrax/_integrate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index 92fccf02..378c5dd2 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -809,7 +809,7 @@ def _save_t1(subsaveat, save_state):
             if subsaveat.t1 or subsaveat.steps:
                 # In this branch we need to replace the last value with tfinal
                 # and yfinal returned by the root finder also if subsaveat.steps
-                # because we deled the last value after the event time above.
+                # because we deleted the last value after the event time above.
                 save_state = _save(tfinal, yfinal, args, subsaveat.fn, save_state)
         return save_state
 

From 0c62f4ce3e91fe632cb866674662e805ff942ea6 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Tue, 18 Jun 2024 17:30:54 +0200
Subject: [PATCH 26/28] Simplified unsaving

---
 diffrax/_integrate.py | 84 +++++++++++--------------------------------
 1 file changed, 21 insertions(+), 63 deletions(-)

diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index 378c5dd2..226a7a7f 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -725,72 +725,30 @@ def _call_real_impl():
         )
 
         # We delete all the saved values after the event time.
-        # For values saved at steps
-        def unsave_step(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState:
-            if subsaveat.steps:
-                save_index = save_state.save_index - 1
-                _ts = save_state.ts.at[save_index].set(jnp.inf)
-                _ys = jtu.tree_map(
-                    lambda _, __ys: __ys.at[save_index].set(jnp.inf),
-                    subsaveat.fn(tfinal, yfinal, args),
-                    save_state.ys,
-                )
-                save_state = eqx.tree_at(
-                    lambda s: s.save_index, save_state, replace_fn=lambda i: i - 1
-                )
-                save_state = SaveState(
-                    saveat_ts_index=save_state.saveat_ts_index,
-                    ts=_ts,
-                    ys=_ys,
-                    save_index=save_index,
-                )
-            return save_state
-
-        save_state = jtu.tree_map(
-            unsave_step,
-            saveat.subs,
-            final_state.save_state,
-            is_leaf=_is_subsaveat,
-        )
-
-        # For values saved at specific times, ts
-        def unsave_ts(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState:
-            if subsaveat.ts is not None:
-                save_state = unsave_ts_impl(subsaveat.ts, subsaveat.fn, save_state)
-            return save_state
-
-        def unsave_ts_impl(ts, fn, save_state: SaveState) -> SaveState:
-            def _cond_fun(_save_state):
-                return (ts[_save_state.saveat_ts_index - 1] > tfinal) & (
-                    _save_state.saveat_ts_index - 1 > 0
-                )
-
-            def _body_fun(_save_state):
-                saveat_ts_index = _save_state.saveat_ts_index - 1
-                _ts = _save_state.ts.at[saveat_ts_index].set(jnp.inf)
-                _ys = jtu.tree_map(
-                    lambda _, __ys: __ys.at[saveat_ts_index].set(jnp.inf),
-                    fn(tfinal, yfinal, args),
-                    _save_state.ys,
-                )
-                return SaveState(
-                    saveat_ts_index=saveat_ts_index,
-                    ts=_ts,
-                    ys=_ys,
-                    save_index=_save_state.save_index - 1,
-                )
-
-            return inner_while_loop(
-                _cond_fun,
-                _body_fun,
-                save_state,
-                max_steps=len(ts),
-                buffers=_inner_buffers,
-                checkpoints=len(ts),
+        def unsave(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState:
+            ts = save_state.ts
+            mask = ts >= tfinal
+            _save_index = save_state.save_index - jnp.sum(mask & (ts < jnp.inf))
+            _saveat_ts_index = save_state.saveat_ts_index - jnp.sum(
+                mask & (ts < jnp.inf)
+            )
+            _ts = jnp.where(mask, jnp.inf, ts)
+            _ys = jtu.tree_map(
+                lambda _, __ys: jnp.where(
+                    mask[(...,) + (jnp.newaxis,) * (__ys.ndim - 1)], jnp.inf, __ys
+                ),
+                subsaveat.fn(tfinal, yfinal, args),
+                save_state.ys,
+            )
+            return SaveState(
+                saveat_ts_index=_saveat_ts_index,
+                ts=_ts,
+                ys=_ys,
+                save_index=_save_index,
             )
 
         save_state = jtu.tree_map(
-            unsave_ts, saveat.subs, save_state, is_leaf=_is_subsaveat
+            unsave, saveat.subs, final_state.save_state, is_leaf=_is_subsaveat
         )
 
         final_state = eqx.tree_at(

From 7588482a3ae8ff8cb13a147816fe96f7d7da2efa Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Tue, 25 Jun 2024 14:18:31 +0200
Subject: [PATCH 27/28] Deleted extra unnecessary argument

---
 diffrax/_integrate.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index 226a7a7f..a3b04028 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -734,10 +734,9 @@ def unsave(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState:
             )
             _ts = jnp.where(mask, jnp.inf, ts)
             _ys = jtu.tree_map(
-                lambda _, __ys: jnp.where(
+                lambda __ys: jnp.where(
                     mask[(...,) + (jnp.newaxis,) * (__ys.ndim - 1)], jnp.inf, __ys
                 ),
-                subsaveat.fn(tfinal, yfinal, args),
                 save_state.ys,
             )
             return SaveState(

From e4935ae036091972201560d1ab11fcb093817ad1 Mon Sep 17 00:00:00 2001
From: cholberg <cholberg24@gmail.com>
Date: Thu, 27 Jun 2024 17:45:48 +0200
Subject: [PATCH 28/28] Changed to strict inequality to be in line with the
 usual saving behviour

---
 diffrax/_integrate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py
index a3b04028..4d08f855 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -727,7 +727,7 @@ def _call_real_impl():
         # We delete all the saved values after the event time.
         def unsave(subsaveat: SubSaveAt, save_state: SaveState) -> SaveState:
             ts = save_state.ts
-            mask = ts >= tfinal
+            mask = ts > tfinal
             _save_index = save_state.save_index - jnp.sum(mask & (ts < jnp.inf))
             _saveat_ts_index = save_state.saveat_ts_index - jnp.sum(
                 mask & (ts < jnp.inf)