Repository: patrick-kidger/diffrax (language: python)

Issue:
Events
Updates to how events are handled in diffrax. The main changes are:

- There is now only one event class, `Event`.
- Multiple `cond_fn` are supported. An event is triggered whenever one of them changes sign.
- Supports differentiable event times if there are real-valued `cond_fn` and a root_finder is provided.
- New tests for the new event implementation.


Some things that still might require a little thinking about:

- Update documentation (and maybe examples, too).
- What to do with outdated tests.
- How to handle the case where the control is not smooth, e.g., Brownian motion.


Events
Updates to how events are handled in diffrax. The main changes are:

- There is now only one event class, `Event`.
- Multiple `cond_fn` are supported. An event is triggered whenever one of them changes sign.
- Supports differentiable event times if there are real-valued `cond_fn` and a root_finder is provided.
- New tests for the new event implementation.


Some things that still might require a little thinking about:

- Update documentation (and maybe examples, too).
- What to do with outdated tests.
- How to handle the case where the control is not smooth, e.g., Brownian motion.




Discussion / hints:
Thank you for the comments, Patrick. I have gone through and made changes accordingly. The only thing I would disagree with is your comment on not having to evaluate `cond_fn`. Correct me if I am wrong, but would we not need to know the sign of `cond_fn` evaluated at the intial condition. I.e., apriori you cannot know if an event happens when the condition function goes from negative to positive or vice versa. Furthermore, if we want `cond_fn` to also possibly take the state as an argument, `event_result` would need to be initialized after `init_state` is defined.
Just letting you know that I've not forgotten about this! I'm trying to focus on getting #344 in first, and then I'm hoping to return to this. They're both quite large changes so I don't want them to step on each other's toes.
Thank you for the comments, Patrick. I have gone through and made changes accordingly. The only thing I would disagree with is your comment on not having to evaluate `cond_fn`. Correct me if I am wrong, but would we not need to know the sign of `cond_fn` evaluated at the intial condition. I.e., apriori you cannot know if an event happens when the condition function goes from negative to positive or vice versa. Furthermore, if we want `cond_fn` to also possibly take the state as an argument, `event_result` would need to be initialized after `init_state` is defined.
Just letting you know that I've not forgotten about this! I'm trying to focus on getting #344 in first, and then I'm hoping to return to this. They're both quite large changes so I don't want them to step on each other's toes.

Fix patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 6efa999d..dec129e5 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -2,11 +2,11 @@ repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.2.2
     hooks:
+      - id: ruff-format  # formatter
+        types_or: [ python, pyi, jupyter ]
       - id: ruff  # linter
         types_or: [ python, pyi, jupyter ]
         args: [ --fix ]
-      - id: ruff-format  # formatter
-        types_or: [ python, pyi, jupyter ]
   - repo: https://github.com/RobertCraigie/pyright-python
     rev: v1.1.350
     hooks:
diff --git a/benchmarks/brownian_tree_times.py b/benchmarks/brownian_tree_times.py
index 66190f45..832e8502 100644
--- a/benchmarks/brownian_tree_times.py
+++ b/benchmarks/brownian_tree_times.py
@@ -1,6 +1,6 @@
 """
 v0.5.0 introduced a new implementation for `diffrax.VirtualBrownianTree` that is
-additionally capable of computing Levy area.
+additionally capable of computing Lévy area.
 
 Here we check the speed of the new implementation against the old implementation, to be
 sure that it is still fast.
diff --git a/diffrax/__init__.py b/diffrax/__init__.py
index e2f094be..67b4ca50 100644
--- a/diffrax/__init__.py
+++ b/diffrax/__init__.py
@@ -22,9 +22,13 @@
     SpaceTimeTimeLevyArea as SpaceTimeTimeLevyArea,
 )
 from ._event import (
-    AbstractDiscreteTerminatingEvent as AbstractDiscreteTerminatingEvent,
-    DiscreteTerminatingEvent as DiscreteTerminatingEvent,
-    SteadyStateEvent as SteadyStateEvent,
+    # 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/_adjoint.py b/diffrax/_adjoint.py
index ff16f921..4ff2dd2c 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,
@@ -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
@@ -563,7 +564,7 @@ def _loop_backsolve_bwd(
     self,
     solver,
     stepsize_controller,
-    discrete_terminating_event,
+    event,
     saveat,
     t0,
     t1,
@@ -573,7 +574,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 +788,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 +830,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 +847,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/_autocitation.py b/diffrax/_autocitation.py
index 484d1ebb..547177ce 100644
--- a/diffrax/_autocitation.py
+++ b/diffrax/_autocitation.py
@@ -10,6 +10,7 @@
 
 from ._adjoint import BacksolveAdjoint, DirectAdjoint, RecursiveCheckpointAdjoint
 from ._brownian import AbstractBrownianPath, VirtualBrownianTree
+from ._custom_types import BrownianIncrement
 from ._heuristics import is_cde, is_sde
 from ._integrate import diffeqsolve
 from ._misc import adjoint_rms_seminorm
@@ -347,28 +348,36 @@ def _virtual_brownian_tree(terms):
     is_vbt = lambda x: isinstance(x, VirtualBrownianTree)
     leaves = jtu.tree_leaves(terms, is_leaf=is_vbt)
     if any(is_vbt(leaf) for leaf in leaves):
-        vbt_ref, _ = _parse_reference_multi(VirtualBrownianTree)
+        vbt_ref, single_seed_ref, _ = _parse_reference_multi(VirtualBrownianTree)
         return (
             r"""
 % You are simulating Brownian motion using a virtual Brownian tree, which was introduced
 % in:
 """
             + vbt_ref
+            + "\n\n"
+            + single_seed_ref
         )
 
 
 @citation_rules.append
 def _space_time_levy_area(terms):
-    has_levy_area = lambda x: isinstance(x, AbstractBrownianPath) and x.levy_area != ""
+    has_levy_area = (
+        lambda x: isinstance(x, AbstractBrownianPath)
+        and x.levy_area != BrownianIncrement
+    )
     leaves = jtu.tree_leaves(terms, is_leaf=has_levy_area)
     if any(has_levy_area(leaf) for leaf in leaves):
-        _, levy_area_ref = _parse_reference_multi(VirtualBrownianTree)
+        _, single_seed_ref, foster_ref = _parse_reference_multi(VirtualBrownianTree)
         return (
             r"""
-% You are simulating Brownian motion using space-time Levy area, the formulae for which
-% were developed in:
-"""
-            + levy_area_ref
+                % You are simulating Brownian motion using Lévy area,
+                % the formulae for which are due to:
+                """
+            + single_seed_ref
+            + "\n\n"
+            + r"""% and Theorem 6.1.6 of:"""
+            + foster_ref
         )
 
 
diff --git a/diffrax/_brownian/base.py b/diffrax/_brownian/base.py
index 1642d315..21618b76 100644
--- a/diffrax/_brownian/base.py
+++ b/diffrax/_brownian/base.py
@@ -41,7 +41,7 @@ def evaluate(
             left-continuous or right-continuous at any jump points, but Brownian
             motion has no jump points.)
         - `use_levy`: If True, the return type will be a `LevyVal`, which contains
-            PyTrees of Brownian increments and their Levy areas.
+            PyTrees of Brownian increments and their Lévy areas.
 
         **Returns:**
 
diff --git a/diffrax/_brownian/path.py b/diffrax/_brownian/path.py
index 1e18f0e5..0333caa5 100644
--- a/diffrax/_brownian/path.py
+++ b/diffrax/_brownian/path.py
@@ -17,6 +17,7 @@
     levy_tree_transpose,
     RealScalarLike,
     SpaceTimeLevyArea,
+    SpaceTimeTimeLevyArea,
 )
 from .._misc import (
     force_bitcast_convert_type,
@@ -45,7 +46,7 @@ class UnsafeBrownianPath(AbstractBrownianPath):
     motion. Hence the restrictions above. (They describe the general case for which the
     correlation structure isn't needed.)
 
-    !!! info "Levy Area"
+    !!! info "Lévy Area"
 
         Can be initialised with `levy_area` set to `diffrax.BrownianIncrement`, or
         `diffrax.SpaceTimeLevyArea`. If `levy_area=diffrax.SpaceTimeLevyArea`, then it
@@ -61,16 +62,18 @@ class UnsafeBrownianPath(AbstractBrownianPath):
     """
 
     shape: PyTree[jax.ShapeDtypeStruct] = eqx.field(static=True)
-    levy_area: type[Union[BrownianIncrement, SpaceTimeLevyArea]] = eqx.field(
-        static=True
-    )
+    levy_area: type[
+        Union[BrownianIncrement, SpaceTimeLevyArea, SpaceTimeTimeLevyArea]
+    ] = eqx.field(static=True)
     key: PRNGKeyArray
 
     def __init__(
         self,
         shape: Union[tuple[int, ...], PyTree[jax.ShapeDtypeStruct]],
         key: PRNGKeyArray,
-        levy_area: type[Union[BrownianIncrement, SpaceTimeLevyArea]],
+        levy_area: type[
+            Union[BrownianIncrement, SpaceTimeLevyArea, SpaceTimeTimeLevyArea]
+        ] = BrownianIncrement,
     ):
         self.shape = (
             jax.ShapeDtypeStruct(shape, lxi.default_floating_dtype())
@@ -129,7 +132,7 @@ def evaluate(
         )
         if use_levy:
             out = levy_tree_transpose(self.shape, out)
-            assert isinstance(out, (BrownianIncrement, SpaceTimeLevyArea))
+            assert isinstance(out, self.levy_area)
         return out
 
     @staticmethod
@@ -138,19 +141,31 @@ def _evaluate_leaf(
         t1: RealScalarLike,
         key,
         shape: jax.ShapeDtypeStruct,
-        levy_area: type[Union[BrownianIncrement, SpaceTimeLevyArea]],
+        levy_area: type[
+            Union[BrownianIncrement, SpaceTimeLevyArea, SpaceTimeTimeLevyArea]
+        ],
         use_levy: bool,
     ):
         w_std = jnp.sqrt(t1 - t0).astype(shape.dtype)
-        w = jr.normal(key, shape.shape, shape.dtype) * w_std
         dt = jnp.asarray(t1 - t0, dtype=complex_to_real_dtype(shape.dtype))
 
-        if levy_area is SpaceTimeLevyArea:
-            key, key_hh = jr.split(key, 2)
+        if levy_area is SpaceTimeTimeLevyArea:
+            key_w, key_hh, key_kk = jr.split(key, 3)
+            w = jr.normal(key_w, shape.shape, shape.dtype) * w_std
+            hh_std = w_std / math.sqrt(12)
+            hh = jr.normal(key_hh, shape.shape, shape.dtype) * hh_std
+            kk_std = w_std / math.sqrt(720)
+            kk = jr.normal(key_kk, shape.shape, shape.dtype) * kk_std
+            levy_val = SpaceTimeTimeLevyArea(dt=dt, W=w, H=hh, K=kk)
+
+        elif levy_area is SpaceTimeLevyArea:
+            key_w, key_hh = jr.split(key, 2)
+            w = jr.normal(key_w, shape.shape, shape.dtype) * w_std
             hh_std = w_std / math.sqrt(12)
             hh = jr.normal(key_hh, shape.shape, shape.dtype) * hh_std
             levy_val = SpaceTimeLevyArea(dt=dt, W=w, H=hh)
         elif levy_area is BrownianIncrement:
+            w = jr.normal(key, shape.shape, shape.dtype) * w_std
             levy_val = BrownianIncrement(dt=dt, W=w)
         else:
             assert False
@@ -168,6 +183,6 @@ def _evaluate_leaf(
     be a tuple of integers, describing the shape of a single JAX array. In that case
     the dtype is chosen to be the default floating-point dtype.
 - `key`: A random key.
-- `levy_area`: Whether to additionally generate Levy area. This is required by some SDE
+- `levy_area`: Whether to additionally generate Lévy area. This is required by some SDE
     solvers.
 """
diff --git a/diffrax/_brownian/tree.py b/diffrax/_brownian/tree.py
index 6488242b..83259567 100644
--- a/diffrax/_brownian/tree.py
+++ b/diffrax/_brownian/tree.py
@@ -21,6 +21,7 @@
     levy_tree_transpose,
     RealScalarLike,
     SpaceTimeLevyArea,
+    SpaceTimeTimeLevyArea,
 )
 from .._misc import (
     is_tuple_of_ints,
@@ -51,9 +52,9 @@
 #   author = {Foster, James M.},
 #   year = {2020}
 # }
-# For more about space-time Levy area see Definition 4.2.1.
-# For the midpoint rule for generating space-time Levy area see Theorem 6.1.6.
-# For the general interpolation rule for space-time Levy area see Theorem 6.1.4.
+# For more about space-time Lévy area see Definition 4.2.1.
+# For the midpoint rule for generating space-time Lévy area see Theorem 6.1.6.
+# For the general interpolation rule for space-time Lévy area see Theorem 6.1.4.
 
 FloatDouble: TypeAlias = tuple[Inexact[Array, " *shape"], Inexact[Array, " *shape"]]
 FloatTriple: TypeAlias = tuple[
@@ -63,6 +64,25 @@
 _BrownianReturn = TypeVar("_BrownianReturn", bound=AbstractBrownianIncrement)
 
 
+# An internal dataclass that holds the rescaled Lévy areas
+# in addition to the non-rescaled ones, for the purposes of
+# taking the difference between two Lévy areas.
+class _LevyVal(eqx.Module):
+    dt: RealScalarLike
+    W: Array
+    H: Optional[Array]
+    bar_H: Optional[Array]
+    K: Optional[Array]
+    bar_K: Optional[Array]
+
+    def __check_init__(self):
+        if self.H is None:
+            assert self.bar_H is None
+            assert self.K is None
+        if self.K is None:
+            assert self.bar_K is None
+
+
 class _State(eqx.Module):
     level: IntScalarLike  # level of the tree
     s: RealScalarLike  # starting time of the interval
@@ -72,60 +92,81 @@ class _State(eqx.Module):
     bkk_s_u_su: Optional[FloatTriple]  # \bar{K}_s, _u, _{s,u}
 
 
-def _levy_diff(_, x0: tuple, x1: tuple) -> AbstractBrownianIncrement:
+def _levy_diff(_, x0: _LevyVal, x1: _LevyVal) -> AbstractBrownianIncrement:
     r"""Computes $(W_{s,u}, H_{s,u})$ from $(W_s, \bar{H}_{s,u})$ and
     $(W_u, \bar{H}_u)$, where $\bar{H}_u = u * H_u$.
 
     **Arguments:**
 
     - `_`: unused, for the purposes of aligning the `jtu.tree_map`.
-    - `x0`: `LevyVal` at time `s`.
-    - `x1`: `LevyVal` at time `u`.
+    - `x0`: `_LevyVal` at time `s`.
+    - `x1`: `_LevyVal` at time `u`.
 
     **Returns:**
 
-    `LevyVal(W_su, H_su)`
+    `AbstractBrownianIncrement(W_su, H_su, K_su)`
     """
+    dtype = jnp.result_type(x0.W)
+    tdtype = complex_to_real_dtype(dtype)
+    su = jnp.asarray(x1.dt - x0.dt, dtype=tdtype)
+    if x0.H is None:  # BM only case
+        assert x1.H is None
+        return BrownianIncrement(dt=su, W=x1.W - x0.W)
+
+    # the following computation is common to the space-time
+    # and the space-time-time Lévy area case
+    assert x0.H is not None
+    assert x1.H is not None
+    assert x0.bar_H is not None
+    assert x1.bar_H is not None
+    w_su = x1.W - x0.W
+    inverse_su = 1 / jnp.where(jnp.abs(su) < jnp.finfo(su).eps, jnp.inf, su)
+    with jax.numpy_dtype_promotion("standard"):
+        u_bb_s = x1.dt * x0.W - x0.dt * x1.W
+        bhh_su = x1.bar_H - x0.bar_H - 0.5 * u_bb_s  # bhh_su = H_{s,u} * (u-s)
+        hh_su = inverse_su * bhh_su
+
+    if x0.K is None:  # space-time Lévy area case
+        return SpaceTimeLevyArea(dt=su, W=w_su, H=hh_su)
 
-    if len(x0) == 2:  # BM only case
-        assert len(x1) == 2
-        dt0, w0 = x0
-        dt1, w1 = x1
-        su = jnp.asarray(dt1 - dt0, dtype=complex_to_real_dtype(w0.dtype))
-        return BrownianIncrement(dt=su, W=w1 - w0)
-
-    elif len(x0) == 4:  # space-time levy area case
-        assert len(x1) == 4
-        dt0, w0, hh0, bhh0 = x0
-        dt1, w1, hh1, bhh1 = x1
-
-        w_su = w1 - w0
-        su = jnp.asarray(dt1 - dt0, dtype=complex_to_real_dtype(w0.dtype))
-        _su = jnp.where(jnp.abs(su) < jnp.finfo(su).eps, jnp.inf, su)
-        inverse_su = 1 / _su
+    elif x0.K is not None:  # space-time-time Lévy area case
+        assert x1.K is not None
+        assert x0.bar_K is not None
+        assert x1.bar_K is not None
         with jax.numpy_dtype_promotion("standard"):
-            u_bb_s = dt1 * w0 - dt0 * w1
-            bhh_su = bhh1 - bhh0 - 0.5 * u_bb_s  # bhh_su = H_{s,u} * (u-s)
-            hh_su = inverse_su * bhh_su
-        return SpaceTimeLevyArea(dt=su, W=w_su, H=hh_su)
+            bkk_su = (
+                x1.bar_K
+                - x0.bar_K
+                - (su / 2) * x0.bar_H
+                + (x0.dt / 2) * bhh_su
+                - ((x1.dt - 2 * x0.dt) / 12) * u_bb_s
+            )
+            su2 = jnp.square(su)
+            inverse_su2 = 1 / jnp.where(jnp.abs(su2) < jnp.finfo(su2).eps, jnp.inf, su2)
+            kk_su = inverse_su2 * bkk_su
+        return SpaceTimeTimeLevyArea(dt=su, W=w_su, H=hh_su, K=kk_su)
     else:
         assert False
 
 
-def _make_levy_val(_, x: tuple) -> AbstractBrownianIncrement:
-    if len(x) == 2:
-        dt, w = x
-        return BrownianIncrement(dt=dt, W=w)
-    elif len(x) == 4:
-        dt, w, hh, bhh = x
-        return SpaceTimeLevyArea(dt=dt, W=w, H=hh)
+def _make_levy_val(_, x: _LevyVal) -> AbstractBrownianIncrement:
+    tdtype = complex_to_real_dtype(x.W)
+    dt = jnp.asarray(x.dt, dtype=tdtype)
+    if x.H is None:
+        return BrownianIncrement(dt=dt, W=x.W)
+    elif x.K is None:
+        return SpaceTimeLevyArea(dt=dt, W=x.W, H=x.H)
     else:
-        assert False
+        return SpaceTimeTimeLevyArea(dt=dt, W=x.W, H=x.H, K=x.K)
 
 
 def _split_interval(
-    pred: BoolScalarLike, x_stu: FloatTriple, x_st_tu: FloatDouble
-) -> FloatTriple:
+    pred: BoolScalarLike, x_stu: Optional[FloatTriple], x_st_tu: Optional[FloatDouble]
+) -> Optional[FloatTriple]:
+    if x_stu is None:
+        assert x_st_tu is None
+        return None
+    assert x_st_tu is not None
     x_s, x_t, x_u = x_stu
     x_st, x_tu = x_st_tu
     x_s = jnp.where(pred, x_t, x_s)
@@ -137,19 +178,16 @@ def _split_interval(
 class VirtualBrownianTree(AbstractBrownianPath):
     """Brownian simulation that discretises the interval `[t0, t1]` to tolerance `tol`.
 
-    !!! info "Levy Area"
+    !!! info "Lévy Area"
+        The parameter `levy_area` can be set to one of:
 
-        Can be initialised with `levy_area` set to `diffrax.BrownianIncrement`, or
-        `diffrax.SpaceTimeLevyArea`. If `levy_area=diffrax.SpaceTimeLevyArea`, then it
-        also computes space-time Lévy area `H`. This is an additional source of
-        randomness required for certain stochastic Runge--Kutta solvers; see
-        [`diffrax.AbstractSRK`][] for more information.
+        - [`diffrax.BrownianIncrement`][] (default, generates the increment of W)
+        - [`diffrax.SpaceTimeLevyArea`][] (generates W and the space-time Lévy area H)
+        - [`diffrax.SpaceTimeTimeLevyArea`][] (generates W, H and the space-time-time
+                                                Lévy area K)
 
-        An error will be thrown during tracing if Lévy area is required but is not
-        available.
-
-        The choice here will impact the Brownian path, so even with the same key, the
-        trajectory will be different depending on the value of `levy_area`.
+        The choice of `levy_area` will impact the Brownian path, so even with the same
+        key, the trajectory will be different depending on the value of `levy_area`.
 
     ??? cite "Reference"
 
@@ -165,7 +203,24 @@ class VirtualBrownianTree(AbstractBrownianPath):
         ```
 
         The implementation here is an improvement on the above, in that it additionally
-        simulates space-time Levy area. This is due to Section 6.1 and Theorem 6.1.6 of
+        simulates space-time and space-time-time Lévy areas, and exactly matches the
+        distribution of the Brownian motion and its Lévy areas at all query times.
+        This is due to the paper
+
+        ```bibtex
+        @misc{jelinčič2024singleseed,
+          title={Single-seed generation of Brownian paths and integrals
+          for adaptive and high order SDE solvers},
+          author={Andraž Jelinčič and James Foster and Patrick Kidger},
+          year={2024},
+          eprint={2405.06464},
+          archivePrefix={arXiv},
+          primaryClass={math.NA}
+        }
+        ```
+
+        and Theorem 6.1.6 of
+
         ```bibtex
         @phdthesis{foster2020a,
           publisher = {University of Oxford},
@@ -175,18 +230,15 @@ class VirtualBrownianTree(AbstractBrownianPath):
           year = {2020}
         }
         ```
-
-        In addition, the implementation here is a further improvement on these by using
-        an interpolation method which ensures the conditional 2nd moments are correct.
     """
 
     t0: RealScalarLike
     t1: RealScalarLike
     tol: RealScalarLike
     shape: PyTree[jax.ShapeDtypeStruct] = eqx.field(static=True)
-    levy_area: type[Union[BrownianIncrement, SpaceTimeLevyArea]] = eqx.field(
-        static=True
-    )
+    levy_area: type[
+        Union[BrownianIncrement, SpaceTimeLevyArea, SpaceTimeTimeLevyArea]
+    ] = eqx.field(static=True)
     key: PyTree[PRNGKeyArray]
     _spline: _Spline = eqx.field(static=True)
 
@@ -199,7 +251,7 @@ def __init__(
         shape: Union[tuple[int, ...], PyTree[jax.ShapeDtypeStruct]],
         key: PRNGKeyArray,
         levy_area: type[
-            Union[BrownianIncrement, SpaceTimeLevyArea]
+            Union[BrownianIncrement, SpaceTimeLevyArea, SpaceTimeTimeLevyArea]
         ] = BrownianIncrement,
         _spline: _Spline = "sqrt",
     ):
@@ -267,12 +319,13 @@ def evaluate(
             # map the interval [self.t0, self.t1] onto [0,1]
             t1 = linear_rescale(self.t0, t1, self.t1)
             levy_1 = self._evaluate(t1)
+            # take the difference between the output for t0 and t1 via Chen's relation
             levy_out = jtu.tree_map(_levy_diff, self.shape, levy_0, levy_1)
 
         levy_out = levy_tree_transpose(self.shape, levy_out)
         # now map [0,1] back onto [self.t0, self.t1]
         levy_out = self._denormalise_bm_inc(levy_out)
-        assert isinstance(levy_out, (BrownianIncrement, SpaceTimeLevyArea))
+        assert isinstance(levy_out, self.levy_area)
         return levy_out if use_levy else levy_out.W
 
     def _evaluate(self, r: RealScalarLike) -> PyTree:
@@ -290,18 +343,24 @@ def _evaluate_leaf(
         key,
         r: RealScalarLike,
         struct: jax.ShapeDtypeStruct,
-    ) -> Union[
-        tuple[RealScalarLike, Array], tuple[RealScalarLike, Array, Array, Array]
-    ]:
+    ) -> _LevyVal:
         shape, dtype = struct.shape, struct.dtype
         tdtype = complex_to_real_dtype(dtype)
-
         t0 = jnp.zeros((), tdtype)
         r = jnp.asarray(r, tdtype)
 
-        if self.levy_area is SpaceTimeLevyArea:
-            state_key, init_key_w, init_key_la = jr.split(key, 3)
-            bhh_1 = jr.normal(init_key_la, shape, dtype) / math.sqrt(12)
+        if self.levy_area is SpaceTimeTimeLevyArea:
+            state_key, init_key_w, init_key_hh, init_key_kk = jr.split(key, 4)
+            bhh_1 = jr.normal(init_key_hh, shape, dtype) / math.sqrt(12)
+            bhh_0 = jnp.zeros_like(bhh_1)
+            bhh = (bhh_0, bhh_1, bhh_1)
+            bkk_1 = jr.normal(init_key_kk, shape, dtype) / math.sqrt(720)
+            bkk_0 = jnp.zeros_like(bkk_1)
+            bkk = (bkk_0, bkk_1, bkk_1)
+
+        elif self.levy_area is SpaceTimeLevyArea:
+            state_key, init_key_w, init_key_hh = jr.split(key, 3)
+            bhh_1 = jr.normal(init_key_hh, shape, dtype) / math.sqrt(12)
             bhh_0 = jnp.zeros_like(bhh_1)
             bhh = (bhh_0, bhh_1, bhh_1)
             bkk = None
@@ -333,7 +392,7 @@ def _body_fun(_state: _State):
             (
                 _t,
                 _w_stu,
-                _w_inc,
+                _w_st_tu,
                 _keys,
                 _bhh_stu,
                 _bhh_st_tu,
@@ -347,15 +406,10 @@ def _body_fun(_state: _State):
             _key_st, _key_tu = _keys
             _key = jnp.where(_cond, _key_st, _key_tu)
 
-            _w = _split_interval(_cond, _w_stu, _w_inc)
-            _bkk = None
-            if self.levy_area is SpaceTimeLevyArea:
-                assert _bhh_stu is not None and _bhh_st_tu is not None
-                _bhh = _split_interval(_cond, _bhh_stu, _bhh_st_tu)
-            elif self.levy_area is BrownianIncrement:
-                _bhh = None
-            else:
-                assert False
+            _w = _split_interval(_cond, _w_stu, _w_st_tu)
+            assert _w is not None
+            _bhh = _split_interval(_cond, _bhh_stu, _bhh_st_tu)
+            _bkk = _split_interval(_cond, _bkk_stu, _bkk_st_tu)
 
             return _State(
                 level=_level,
@@ -369,14 +423,119 @@ def _body_fun(_state: _State):
         final_state = lax.while_loop(_cond_fun, _body_fun, init_state)
 
         s = final_state.s
-        su = 2.0**-final_state.level
+        su = jnp.asarray(2.0**-final_state.level, dtype=tdtype)
 
         sr = jax.nn.relu(r - s)
         # make sure su = sr + ru regardless of cancellation error
         ru = jax.nn.relu(su - sr)
 
         w_s, w_u, w_su = final_state.w_s_u_su
-        if self.levy_area is SpaceTimeLevyArea:
+
+        if self.levy_area is SpaceTimeTimeLevyArea:
+            # Based on Theorem 3.7 from the paper on
+            # Single-seed generation of Brownian paths and integrals
+            assert final_state.bhh_s_u_su is not None
+            assert final_state.bkk_s_u_su is not None
+            bhh_s, bhh_u, bhh_su = final_state.bhh_s_u_su
+            bkk_s, bkk_u, bkk_su = final_state.bkk_s_u_su
+
+            su3 = jnp.power(su, 3)
+            sr_by_su = sr / su
+            sr_by_su_3 = jnp.power(sr_by_su, 3)
+            sr_by_su_5 = jnp.power(sr_by_su, 5)
+            ru_by_su = ru / su
+            sr_ru_by_su2 = sr_by_su * ru_by_su
+            sr2 = jnp.square(sr)
+            ru2 = jnp.square(ru)
+            su2 = jnp.square(su)
+
+            # compute the mean of (W_sr, H_sr, K_sr) conditioned on
+            # (W_s, H_s, K_s, W_u, H_u, K_u)
+            with jax.numpy_dtype_promotion("standard"):
+                bb_mean = (6 * sr_ru_by_su2 / su) * bhh_su + (
+                    120 * sr_ru_by_su2 * (0.5 - sr_by_su) / su2
+                ) * bkk_su
+                w_mean = sr_by_su * w_su + bb_mean
+                h_mean = (sr_by_su**2 / su) * (bhh_su + (30 * ru_by_su / su) * bkk_su)
+                k_mean = (sr_by_su_3 / su2) * bkk_su
+
+                # compute the covariance matrix of (W_sr, H_sr, K_sr) conditioned on
+                # (W_s, H_s, K_s, W_u, H_u, K_u)
+                ww_cov = (
+                    sr_by_su * ru_by_su * ((sr - ru) ** 4 + 4 * (sr2 * ru2))
+                ) / su3
+                wh_cov = -(sr_by_su_3 * ru_by_su * (sr2 - 3 * sr * ru + 6 * ru2)) / (
+                    2 * su
+                )
+                wk_cov = (sr_by_su**4) * ru_by_su * (sr - ru) / 12
+                hh_cov = (sr / 12) * (
+                    1 - sr_by_su_3 * (sr2 + 2 * sr * ru + 16 * ru2) / su2
+                )
+                hk_cov = -(ru / 24) * sr_by_su_5
+                kk_cov = (sr / 720) * (1.0 - sr_by_su_5)
+
+            cov = jnp.array(
+                [
+                    [ww_cov, wh_cov, wk_cov],
+                    [wh_cov, hh_cov, hk_cov],
+                    [wk_cov, hk_cov, kk_cov],
+                ]
+            )
+
+            if self._spline == "sqrt":
+                # NOTE: jr.multivariate_normal is not compatible with jnp.float16,
+                # so we need to cast to jnp.float32 before calling it.
+                with jax.numpy_dtype_promotion("standard"):
+                    dtype_atleast32 = jnp.result_type(dtype, jnp.float32)
+                cov = jnp.asarray(cov, dtype_atleast32)
+                hat_y = jr.multivariate_normal(
+                    final_state.key,
+                    jnp.zeros((3,), dtype_atleast32),
+                    cov,
+                    shape=shape,
+                    dtype=dtype_atleast32,
+                    method="svd",
+                )
+                hat_y = jnp.asarray(hat_y, dtype)
+
+            elif self._spline == "zero":
+                hat_y = jnp.zeros(shape=shape + (3,), dtype=dtype)
+            else:
+                raise ValueError(
+                    f"When levy_area='space-time-time', only 'sqrt' and"
+                    f" 'zero' splines are permitted, got {self._spline}."
+                )
+
+            hat_w_sr, hat_hh_sr, hat_kk_sr = [
+                x.squeeze(axis=-1) for x in jnp.split(hat_y, 3, axis=-1)
+            ]
+            assert hat_w_sr.shape == hat_hh_sr.shape == hat_kk_sr.shape == shape
+
+            w_sr = w_mean + hat_w_sr
+            w_r = w_s + w_sr
+
+            with jax.numpy_dtype_promotion("standard"):
+                r_bb_s = r * w_s - s * w_r
+
+                bhh_sr = sr * (h_mean + hat_hh_sr)
+                bhh_r = bhh_s + bhh_sr + 0.5 * r_bb_s
+
+                bkk_sr = sr2 * (k_mean + hat_kk_sr)
+                bkk_r = (
+                    bkk_s
+                    + bkk_sr
+                    + (sr / 2) * bhh_s
+                    - (s / 2) * bhh_sr
+                    + ((r - 2 * s) / 12) * r_bb_s
+                )
+
+                inverse_r = 1 / jnp.where(jnp.square(r) < jnp.finfo(r).eps, jnp.inf, r)
+                hh_r = inverse_r * bhh_r
+                kk_r = inverse_r**2 * bkk_r
+
+            return _LevyVal(dt=r, W=w_r, H=hh_r, bar_H=bhh_r, K=kk_r, bar_K=bkk_r)
+
+        elif self.levy_area is SpaceTimeLevyArea:
             # This is based on Theorem 6.1.4 of Foster's thesis (see above).
 
             assert final_state.bhh_s_u_su is not None
@@ -418,6 +577,8 @@ def _body_fun(_state: _State):
                 inverse_r = 1 / jnp.where(jnp.abs(r) < jnp.finfo(r).eps, jnp.inf, r)
                 hh_r = inverse_r * bhh_r
 
+            return _LevyVal(dt=r, W=w_r, H=hh_r, bar_H=bhh_r, K=None, bar_K=None)
+
         elif self.levy_area is BrownianIncrement:
             with jax.numpy_dtype_promotion("standard"):
                 w_mean = w_s + sr / su * w_su
@@ -432,13 +593,11 @@ def _body_fun(_state: _State):
                 else:
                     assert False
             w_r = w_mean + bb
-            return r, w_r
+            return _LevyVal(dt=r, W=w_r, H=None, bar_H=None, K=None, bar_K=None)
 
         else:
             assert False
 
-        return r, w_r, hh_r, bhh_r
-
     def _brownian_arch(
         self, _state: _State, shape, dtype
     ) -> tuple[
@@ -458,7 +617,7 @@ def _brownian_arch(
          and also returns `w_st` and `w_tu` in addition to just `w_t`. Same for `bhh`
          if it is not None.
          Note that the inputs and outputs already contain `bkk`. These values are
-         there for the sake of a future extension with "space-time-time" Levy area
+         there for the sake of a future extension with "space-time-time" Lévy area
          and should be None for now.
 
         **Arguments:**
@@ -477,61 +636,113 @@ def _brownian_arch(
         - `bhh_st_tu`: (optional) $(\bar{H}_{s,t}, \bar{H}_{t,u})$
         - `bkk_stu`: (optional) $(\bar{K}_s, \bar{K}_t, \bar{K}_u)$
         - `bkk_st_tu`: (optional) $(\bar{K}_{s,t}, \bar{K}_{t,u})$
+
         """
         key_st, midpoint_key, key_tu = jr.split(_state.key, 3)
         keys = (key_st, key_tu)
-        su = 2.0**-_state.level
+        tdtype = complex_to_real_dtype(dtype)
+        su = jnp.asarray(2.0**-_state.level, dtype=tdtype)
         st = su / 2
-        s = _state.s
+        s = jnp.asarray(_state.s, dtype=tdtype)
         t = s + st
         root_su = jnp.sqrt(su)
 
         w_s, w_u, w_su = _state.w_s_u_su
 
-        if self.levy_area is SpaceTimeLevyArea:
-            assert _state.bhh_s_u_su is not None
-            assert _state.bkk_s_u_su is None
-            bhh_s, bhh_u, bhh_su = _state.bhh_s_u_su
-
-            z1_key, z2_key = jr.split(midpoint_key, 2)
-            z1 = jr.normal(z1_key, shape, dtype)
-            z2 = jr.normal(z2_key, shape, dtype)
-            z = z1 * (root_su / 4)
-            n = z2 * jnp.sqrt(su / 12)
-
-            w_term1 = w_su / 2
-            w_term2 = 3 / (2 * su) * bhh_su + z
-            w_st = w_term1 + w_term2
-            w_tu = w_term1 - w_term2
-            w_st_tu = (w_st, w_tu)
-
-            bhh_term1 = bhh_su / 8 - su / 4 * z
-            bhh_term2 = su / 4 * n
-            bhh_st = bhh_term1 + bhh_term2
-            bhh_tu = bhh_term1 - bhh_term2
-            bhh_st_tu = (bhh_st, bhh_tu)
-
-            w_t = w_s + w_st
-            w_stu = (w_s, w_t, w_u)
-            with jax.numpy_dtype_promotion("standard"):
+        with jax.numpy_dtype_promotion("standard"):
+            if self.levy_area is SpaceTimeTimeLevyArea:
+                assert _state.bhh_s_u_su is not None
+                assert _state.bkk_s_u_su is not None
+
+                bhh_s, bhh_u, bhh_su = _state.bhh_s_u_su
+                bkk_s, bkk_u, bkk_su = _state.bkk_s_u_su
+
+                z1_key, z2_key, z3_key = jr.split(midpoint_key, 3)
+                z1 = jr.normal(z1_key, shape, dtype)
+                z2 = jr.normal(z2_key, shape, dtype)
+                z3 = jr.normal(z3_key, shape, dtype)
+
+                z = z1 * jnp.sqrt(su / 16)
+                x1 = z2 * jnp.sqrt(su / 768)
+                x2 = z3 * jnp.sqrt(su / 2880)
+
+                su2 = su**2
+
+                w_term1 = w_su / 2
+                w_term2 = (3 / (2 * su)) * bhh_su + z
+                w_st = w_term1 + w_term2
+                w_tu = w_term1 - w_term2
+                bhh_term1 = bhh_su / 8 - (st / 2) * z
+                bhh_term2 = (15 / (8 * su)) * bkk_su + st * x1
+                bhh_st = bhh_term1 + bhh_term2
+                bhh_tu = bhh_term1 - bhh_term2
+                bkk_term1 = bkk_su / 32 - (su2 / 8) * x1
+                bkk_term2 = (su2 / 4) * x2
+                bkk_st = bkk_term1 + bkk_term2
+                bkk_tu = bkk_term1 - bkk_term2
+                w_st_tu = (w_st, w_tu)
+                bhh_st_tu = (bhh_st, bhh_tu)
+                bkk_st_tu = (bkk_st, bkk_tu)
+
+                w_t = w_s + w_st
+                t_bb_s = t * w_s - s * w_t
+                bhh_t = bhh_s + bhh_st + t_bb_s / 2
+                bkk_t = (
+                    bkk_s
+                    + bkk_st
+                    + (st / 2) * bhh_s
+                    - (s / 2) * bhh_st
+                    + ((t - 2 * s) / 12) * t_bb_s
+                )
+
+                w_stu = (w_s, w_t, w_u)
+                bhh_stu = (bhh_s, bhh_t, bhh_u)
+                bkk_stu = (bkk_s, bkk_t, bkk_u)
+
+            elif self.levy_area is SpaceTimeLevyArea:
+                assert _state.bhh_s_u_su is not None
+                assert _state.bkk_s_u_su is None
+                bhh_s, bhh_u, bhh_su = _state.bhh_s_u_su
+
+                z1_key, z2_key = jr.split(midpoint_key, 2)
+                z1 = jr.normal(z1_key, shape, dtype)
+                z2 = jr.normal(z2_key, shape, dtype)
+                z = z1 * (root_su / 4)
+                n = z2 * jnp.sqrt(su / 12)
+
+                w_term1 = w_su / 2
+                w_term2 = (3 / (2 * su)) * bhh_su + z
+                w_st = w_term1 + w_term2
+                w_tu = w_term1 - w_term2
+                w_st_tu = (w_st, w_tu)
+
+                bhh_term1 = bhh_su / 8 - su / 4 * z
+                bhh_term2 = (su / 4) * n
+                bhh_st = bhh_term1 + bhh_term2
+                bhh_tu = bhh_term1 - bhh_term2
+                bhh_st_tu = (bhh_st, bhh_tu)
+
+                w_t = w_s + w_st
+                w_stu = (w_s, w_t, w_u)
+
                 bhh_t = bhh_s + bhh_st + 0.5 * (t * w_s - s * w_t)
-            bhh_stu = (bhh_s, bhh_t, bhh_u)
-            bkk_stu = None
-            bkk_st_tu = None
+                bhh_stu = (bhh_s, bhh_t, bhh_u)
+                bkk_stu = None
+                bkk_st_tu = None
 
-        elif self.levy_area is BrownianIncrement:
-            assert _state.bhh_s_u_su is None
-            assert _state.bkk_s_u_su is None
-            mean = 0.5 * w_su
-            w_term2 = root_su / 2 * jr.normal(midpoint_key, shape, dtype)
-            w_st = mean + w_term2
-            w_tu = mean - w_term2
-            w_st_tu = (w_st, w_tu)
-            w_t = w_s + w_st
-            w_stu = (w_s, w_t, w_u)
-            bhh_stu, bhh_st_tu, bkk_stu, bkk_st_tu = None, None, None, None
+            elif self.levy_area is BrownianIncrement:
+                assert _state.bhh_s_u_su is None
+                assert _state.bkk_s_u_su is None
+                mean = 0.5 * w_su
+                w_term2 = (root_su / 2) * jr.normal(midpoint_key, shape, dtype)
+                w_st = mean + w_term2
+                w_tu = mean - w_term2
+                w_st_tu = (w_st, w_tu)
+                w_t = w_s + w_st
+                w_stu = (w_s, w_t, w_u)
+                bhh_stu, bhh_st_tu, bkk_stu, bkk_st_tu = None, None, None, None
 
-        else:
-            assert False
+            else:
+                assert False
 
         return t, w_stu, w_st_tu, keys, bhh_stu, bhh_st_tu, bkk_stu, bkk_st_tu
diff --git a/diffrax/_custom_types.py b/diffrax/_custom_types.py
index 69853e98..44e9258f 100644
--- a/diffrax/_custom_types.py
+++ b/diffrax/_custom_types.py
@@ -47,6 +47,8 @@
 Control = PyTree[Shaped[ArrayLike, "?*control"], "C"]
 Args = PyTree[Any]
 
+BM = PyTree[Shaped[ArrayLike, "?*bm"], "BM"]
+
 DenseInfo = dict[str, PyTree[Array]]
 DenseInfos = dict[str, PyTree[Shaped[Array, "times-1 ..."]]]
 BufferDenseInfos = dict[str, PyTree[eqxi.MaybeBuffer[Shaped[Array, "times ..."]]]]
@@ -54,34 +56,34 @@
 
 
 class AbstractBrownianIncrement(eqx.Module):
-    dt: eqx.AbstractVar[PyTree[FloatScalarLike]]
-    W: eqx.AbstractVar[PyTree[Array]]
+    dt: eqx.AbstractVar[PyTree[FloatScalarLike, "BM"]]
+    W: eqx.AbstractVar[BM]
 
 
 class AbstractSpaceTimeLevyArea(AbstractBrownianIncrement):
-    H: eqx.AbstractVar[PyTree[Array]]
+    H: eqx.AbstractVar[BM]
 
 
 class AbstractSpaceTimeTimeLevyArea(AbstractSpaceTimeLevyArea):
-    K: eqx.AbstractVar[PyTree[Array]]
+    K: eqx.AbstractVar[BM]
 
 
 class BrownianIncrement(AbstractBrownianIncrement):
-    dt: PyTree[FloatScalarLike]
-    W: PyTree[Array]
+    dt: PyTree[FloatScalarLike, "BM"]
+    W: BM
 
 
 class SpaceTimeLevyArea(AbstractSpaceTimeLevyArea):
-    dt: PyTree[FloatScalarLike]
-    W: PyTree[Array]
-    H: PyTree[Array]
+    dt: PyTree[FloatScalarLike, "BM"]
+    W: BM
+    H: BM
 
 
 class SpaceTimeTimeLevyArea(AbstractSpaceTimeTimeLevyArea):
-    dt: PyTree[FloatScalarLike]
-    W: PyTree[Array]
-    H: PyTree[Array]
-    K: PyTree[Array]
+    dt: PyTree[FloatScalarLike, "BM"]
+    W: BM
+    H: BM
+    K: BM
 
 
 def levy_tree_transpose(
diff --git a/diffrax/_event.py b/diffrax/_event.py
index a596d8ed..3168b297 100644
--- a/diffrax/_event.py
+++ b/diffrax/_event.py
@@ -1,98 +1,187 @@
 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 ._custom_types import BoolScalarLike, RealScalarLike
+from ._custom_types import BoolScalarLike, FloatScalarLike, 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**
+class Event(eqx.Module):
+    """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.
 
-        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]
+    Instances of this class should be passed as the `event` argument of
+    [`diffrax.diffeqsolve`][].
+    """
 
-    def __call__(self, state, **kwargs):
-        return self.cond_fn(state, **kwargs)
+    cond_fn: PyTree[Callable[..., Union[BoolScalarLike, RealScalarLike]]]
+    root_finder: Optional[optx.AbstractRootFinder] = None
+
+
+Event.__init__.__doc__ = """**Arguments:**
+
+- `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](../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
+
+    $\\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 vector_field(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)
+    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...
+    ```
+"""
 
 
-DiscreteTerminatingEvent.__init__.__doc__ = """**Arguments:**
+def steady_state_event(
+    rtol: Optional[float] = None,
+    atol: Optional[float] = None,
+    norm: Optional[Callable[[PyTree[Array]], RealScalarLike]] = None,
+):
+    """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`][].
 
-- `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.
-"""
+    **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.
 
-class SteadyStateEvent(AbstractDiscreteTerminatingEvent):
-    """Terminates the solve once it reaches a steady state."""
+    **Returns:**
 
-    rtol: Optional[float] = None
-    atol: Optional[float] = None
-    norm: Callable[[PyTree[Array]], RealScalarLike] = optx.rms_norm
+    A function `f(t, y, args, **kwargs)`, that can be passed to
+    `diffrax.Event(cond_fn=..., ...)`.
+    """
 
-    def __call__(self, state, *, terms, args, solver, stepsize_controller, **kwargs):
+    def _cond_fn(t, y, args, *, terms, 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."
+            "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 self.rtol is None:
+        if 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:
+            _rtol = rtol
+        if atol is None:
             if isinstance(stepsize_controller, AbstractAdaptiveStepSizeController):
                 _atol = stepsize_controller.atol
             else:
                 raise ValueError(msg)
         else:
-            _atol = self.atol
+            _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, state.tprev, state.y, args)
-        return self.norm(vf) < _atol + _rtol * self.norm(state.y)
+        vf = solver.func(terms, t, y, args)
+        return _norm(vf) < _atol + _rtol * _norm(y)
 
+    return _cond_fn
 
-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.
-"""
+#
+# 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)(
+            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 6cb25140..615eafe9 100644
--- a/diffrax/_integrate.py
+++ b/diffrax/_integrate.py
@@ -9,26 +9,34 @@
     Optional,
     Tuple,
     TYPE_CHECKING,
+    Union,
 )
 
 import equinox as eqx
 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
+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 (
+    AbstractDiscreteTerminatingEvent,
+    DiscreteTerminatingEventToCondFn,
+    Event,
+)
 from ._global_interpolation import DenseInterpolation
 from ._heuristics import is_sde, is_unsafe_sde
 from ._misc import linear_rescale, static_select
@@ -67,23 +75,41 @@ 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]
+    #
+    # 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]]
 
 
 def _is_none(x: Any) -> bool:
@@ -243,7 +269,7 @@ def loop(
     *,
     solver,
     stepsize_controller,
-    discrete_terminating_event,
+    event,
     saveat,
     t0,
     t1,
@@ -461,6 +487,87 @@ 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
+
+            def _outer_cond_fn(cond_fn_i, old_event_value_i):
+                new_event_value_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,
+                )
+                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 "
+                        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:
+                    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(
+                _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,
+                jtu.tree_structure((0, 0)),
+                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)
+            result = RESULTS.where(
+                had_event,
+                RESULTS.event_occurred,
+                result,
+            )
+
         new_state = State(
             y=y,
             tprev=tprev,
@@ -473,32 +580,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,
+            event_tprev=event_tprev,
+            event_tnext=event_tnext,
+            event_dense_info=event_dense_info,
+            event_values=event_values,
+            event_mask=event_mask,
         )
 
-        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,
-            )
-            result = RESULTS.where(
-                discrete_terminating_event_occurred,
-                RESULTS.discrete_terminating_event_occurred,
-                result,
-            )
-            new_state = eqx.tree_at(lambda s: s.result, new_state, result)
-
         return (
             new_state,
             (type(new_state.made_jump) is not bool),
@@ -529,16 +621,153 @@ 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
 
-    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
+    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 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)
+                    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),
+        )
+
+        # 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
+            _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
+                ),
+                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, saveat.subs, final_state.save_state, is_leaf=_is_subsaveat
+        )
+
+        final_state = eqx.tree_at(
+            lambda s: s.save_state,
+            final_state,
+            save_state,
+            is_leaf=_is_none,
+        )
+
+    def _save_t1(subsaveat, save_state):
+        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)
+        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
+                # because we deleted the last value after the event time above.
+                save_state = _save(tfinal, yfinal, args, subsaveat.fn, save_state)
         return save_state
 
     save_state = jtu.tree_map(
@@ -549,10 +778,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
 
 
@@ -566,6 +793,7 @@ class SaveAt(eqx.Module):  # noqa: F811
 
 
 @eqx.filter_jit
+@eqxi.doc_remove_args("discrete_terminating_event")
 def diffeqsolve(
     terms: PyTree[AbstractTerm],
     solver: AbstractSolver,
@@ -578,13 +806,15 @@ 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(),
     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.
 
@@ -629,8 +859,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.
@@ -640,9 +870,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.
@@ -693,6 +921,24 @@ 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.",
+            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 = (
@@ -952,25 +1198,20 @@ 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
@@ -980,6 +1221,91 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
     # Progress meter
     progress_meter_state = progress_meter.init()
 
+    # Events
+    if event is None:
+        event_tprev = None
+        event_tnext = None
+        event_dense_info = None
+        event_values = None
+        event_mask = None
+    else:
+        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]
+        )
+
+        def _outer_cond_fn(cond_fn_i):
+            event_value_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,
+            )
+            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 "
+                    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:
+                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(
+            _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,
+            jtu.tree_structure((0, 0)),
+            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)
+        result = RESULTS.where(
+            had_event,
+            RESULTS.event_occurred,
+            result,
+        )
+        del had_event, event_structure, event_mask_leaves, event_values__mask
+
     # Initialise state
     init_state = State(
         y=y0,
@@ -997,6 +1323,11 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
         dense_infos=dense_infos,
         dense_save_index=dense_save_index,
         progress_meter_state=progress_meter_state,
+        event_tprev=event_tprev,
+        event_tnext=event_tnext,
+        event_dense_info=event_dense_info,
+        event_values=event_values,
+        event_mask=event_mask,
     )
 
     #
@@ -1008,7 +1339,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,9 +1362,10 @@ 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)
+
     # 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
@@ -1069,8 +1401,10 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
         "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,
@@ -1082,6 +1416,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..f1b8d21b 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,33 @@
 
 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."
     )
+    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.",
+        stacklevel=2,
+    )
+    return self.event_occurred
+
+
+RESULTS.discrete_terminating_event_occurred = discrete_terminating_event_occurred  # pyright: ignore[reportAttributeAccessIssue]
 
 
 def is_okay(result: RESULTS) -> Bool[Array, ""]:
@@ -32,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.discrete_terminating_event_occurred
+    return result == RESULTS.event_occurred
 
 
 def update_result(old_result: RESULTS, new_result: RESULTS) -> RESULTS:
@@ -67,12 +83,17 @@ 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.RESULTS`][] 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.
     - `made_jump`: If saved, the final internal state for the jump tracker.
+    - `event_mask`: If using [events](../events), a boolean mask indicating which event
+        triggered. This is a PyTree of bools, with the same PyTree stucture as the event
+        condition functions. It will be all `False` if no events triggered; otherwise it
+        will have precisely one `True`, corresponding to the event that triggered.
 
     !!! note
 
@@ -99,6 +120,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/diffrax/_solver/srk.py b/diffrax/_solver/srk.py
index 494ceacd..8ac53715 100644
--- a/diffrax/_solver/srk.py
+++ b/diffrax/_solver/srk.py
@@ -143,8 +143,8 @@ def __post_init__(self):
             assert self.coeffs_hh.check() == num_stages
         if self.coeffs_kk is not None:
             assert self.coeffs_hh is not None, (
-                "If space-time-time Levy area (K) is used,"
-                " space-time Levy area (H) must also be used."
+                "If space-time-time Lévy area (K) is used,"
+                " space-time Lévy area (H) must also be used."
             )
             assert type(self.coeffs_kk) is type(self.coeffs_w)
             assert self.coeffs_kk.check() == num_stages
@@ -189,9 +189,9 @@ def __post_init__(self):
 - `coeffs_w`: An instance of `AdditiveCoeffs` or `GeneralCoeffs`, providing the
     coefficients of the Brownian motion increments.
 - `coeffs_hh`: An instance of `AdditiveCoeffs` or `GeneralCoeffs`, providing the
-    coefficients of the space-time Levy area.
+    coefficients of the space-time Lévy area.
 - `coeffs_kk`: An instance of `AdditiveCoeffs` or `GeneralCoeffs`, providing the
-    coefficients of the space-time-time Levy area.
+    coefficients of the space-time-time Lévy area.
 - `ignore_stage_f`: Optional. A NumPy array of length `s` of booleans. If `True` at
     stage `j`, the vector field of the drift term will not be evaluated at stage `j`.
 - `ignore_stage_g`: Optional. A NumPy array of length `s` of booleans. If `True` at
@@ -203,12 +203,11 @@ class AbstractSRK(AbstractSolver[_SolverState]):
     r"""A general Stochastic Runge-Kutta method.
 
     This accepts `terms` of the form
-    `MultiTerm(ODETerm(drift), ControlTerm(diffusion, brownian_motion))` or
-    `MultiTerm(ODETerm(drift), WeaklyDiagonalControlTerm(diffusion, brownian_motion))`.
+    `MultiTerm(ODETerm(drift), ControlTerm(diffusion, brownian_motion))`.
     Depending on the solver, the Brownian motion might need to generate
-    different types of Levy areas, specified by the `minimal_levy_area` attribute.
+    different types of Lévy areas, specified by the `minimal_levy_area` attribute.
 
-    For example, the [`diffrax.ShARK`][] solver requires space-time Levy area, so
+    For example, the [`diffrax.ShARK`][] solver requires space-time Lévy area, so
     it will have `minimal_levy_area = AbstractSpaceTimeLevyArea` and the Brownian
     motion must be initialised with `levy_area=SpaceTimeLevyArea`.
 
@@ -260,10 +259,10 @@ class AbstractSRK(AbstractSolver[_SolverState]):
     term_compatible_contr_kwargs = (dict(), dict(use_levy=True))
     tableau: AbstractClassVar[StochasticButcherTableau]
 
-    # Indicates the type of Levy area used by the solver.
-    # The BM must generate at least this type of Levy area, but can generate
-    # more. E.g. if the solver uses space-time Levy area, then the BM generates
-    # space-time-time Levy area as well that is fine. The other way around would
+    # Indicates the type of Lévy area used by the solver.
+    # The BM must generate at least this type of Lévy area, but can generate
+    # more. E.g. if the solver uses space-time Lévy area, then the BM generates
+    # space-time-time Lévy area as well that is fine. The other way around would
     # not work. This is mostly an easily readable indicator so that methods know
     # what kind of BM to use.
     @property
@@ -288,7 +287,7 @@ def init(
         args: PyTree,
     ) -> _SolverState:
         del t1
-        # Check that the diffusion has the correct Levy area
+        # Check that the diffusion has the correct Lévy area
         _, diffusion = terms.terms
 
         if self.tableau.is_additive_noise():
@@ -372,12 +371,12 @@ def make_zeros_aux(leaf):
         b_levy_list = []
 
         levy_areas = []
-        if self.tableau.coeffs_hh is not None:  # space-time Levy area
+        if self.tableau.coeffs_hh is not None:  # space-time Lévy area
             assert isinstance(bm_inc, AbstractSpaceTimeLevyArea)
             levy_areas.append(bm_inc.H)
             b_levy_list.append(jnp.asarray(self.tableau.coeffs_hh.b_sol, dtype=dtype))
 
-            if self.tableau.coeffs_kk is not None:  # space-time-time Levy area
+            if self.tableau.coeffs_kk is not None:  # space-time-time Lévy area
                 assert isinstance(bm_inc, AbstractSpaceTimeTimeLevyArea)
                 levy_areas.append(bm_inc.K)
                 b_levy_list.append(
@@ -397,7 +396,7 @@ def aux_add_levy(w_leaf, *levy_leaves):
 
         levylist_kgs = []  # will contain levy * g(t0 + c_j * h, z_j) for each stage j
         # where levy is either H or K (if those entries exist)
-        # this is similar to h_kfs or w_kgs, but for the Levy area(s)
+        # this is similar to h_kfs or w_kgs, but for the Lévy area(s)
 
         if self.tableau.is_additive_noise():  # additive noise
             # compute g once since it is constant
@@ -413,12 +412,12 @@ def _comp_g(_t):
             w_kgs = diffusion.prod(g0, w)
             a_w = jnp.asarray(self.tableau.coeffs_w.a, dtype=dtype)
 
-            if self.tableau.coeffs_hh is not None:  # space-time Levy area
+            if self.tableau.coeffs_hh is not None:  # space-time Lévy area
                 assert isinstance(bm_inc, AbstractSpaceTimeLevyArea)
                 levylist_kgs.append(diffusion.prod(g0, bm_inc.H))
                 a_levy.append(jnp.asarray(self.tableau.coeffs_hh.a, dtype=dtype))
 
-            if self.tableau.coeffs_kk is not None:  # space-time-time Levy area
+            if self.tableau.coeffs_kk is not None:  # space-time-time Lévy area
                 assert isinstance(bm_inc, AbstractSpaceTimeTimeLevyArea)
                 levylist_kgs.append(diffusion.prod(g0, bm_inc.K))
                 a_levy.append(jnp.asarray(self.tableau.coeffs_kk.a, dtype=dtype))
@@ -436,11 +435,11 @@ def _comp_g(_t):
             w_kgs = make_zeros()
             a_w = self._embed_a_lower(self.tableau.coeffs_w.a, dtype)
 
-            # do the same for each type of Levy area
-            if self.tableau.coeffs_hh is not None:  # space-time Levy area
+            # do the same for each type of Lévy area
+            if self.tableau.coeffs_hh is not None:  # space-time Lévy area
                 levylist_kgs.append(make_zeros())
                 a_levy.append(self._embed_a_lower(self.tableau.coeffs_hh.a, dtype))
-            if self.tableau.coeffs_kk is not None:  # space-time-time Levy area
+            if self.tableau.coeffs_kk is not None:  # space-time-time Lévy area
                 levylist_kgs.append(make_zeros())
                 a_levy.append(self._embed_a_lower(self.tableau.coeffs_kk.a, dtype))
 
@@ -581,7 +580,7 @@ def compute_and_insert_kg_j(_w_kgs_in, _levylist_kgs_in):
             # In the additive noise case (i.e. when g is independent of y),
             # we still need a correction term in case the diffusion vector field
             # g depends on t. This term is of the form $(g1 - g0) * (0.5*W_n - H_n)$.
-            if self.tableau.coeffs_hh is not None:  # space-time Levy area
+            if self.tableau.coeffs_hh is not None:  # space-time Lévy area
                 assert isinstance(bm_inc, AbstractSpaceTimeLevyArea)
                 time_var_contr = (bm_inc.W**ω - 2.0 * bm_inc.H**ω).ω
                 time_var_term = diffusion.prod(g_delta, time_var_contr)
diff --git a/diffrax/_step_size_controller/adaptive.py b/diffrax/_step_size_controller/adaptive.py
index eb14d712..9d181c95 100644
--- a/diffrax/_step_size_controller/adaptive.py
+++ b/diffrax/_step_size_controller/adaptive.py
@@ -610,7 +610,7 @@ def _scale(_y0, _y1_candidate, _y_error):
         # a grad API boundary as part of a larger model.)
         factor = lax.stop_gradient(factor)
         factor = eqxi.nondifferentiable(factor)
-        dt = prev_dt * factor.astype(prev_dt)
+        dt = prev_dt * factor.astype(jnp.result_type(prev_dt))
 
         # E.g. we failed an implicit step, so y_error=inf, so inv_scaled_error=0,
         # so factor=factormin, and we shrunk our step.
diff --git a/diffrax/_term.py b/diffrax/_term.py
index 4fe4e9bc..8d031b95 100644
--- a/diffrax/_term.py
+++ b/diffrax/_term.py
@@ -1,5 +1,6 @@
 import abc
 import operator
+import warnings
 from collections.abc import Callable
 from typing import cast, Generic, Optional, TypeVar, Union
 
@@ -7,6 +8,7 @@
 import jax
 import jax.numpy as jnp
 import jax.tree_util as jtu
+import lineax as lx
 import numpy as np
 from equinox.internal import ω
 from jaxtyping import ArrayLike, PyTree, PyTreeDef
@@ -82,6 +84,10 @@ def prod(self, vf: _VF, control: _Control) -> Y:
         control $x$, this computes $f(t, y(t), args) \Delta x(t)$ given
         $f(t, y(t), args)$ and $\Delta x(t)$.
 
+        !!! note
+
+            This function must be bilinear.
+
         **Arguments:**
 
         - `vf`: The vector field evaluation; a PyTree of structure $S$.
@@ -91,10 +97,6 @@ def prod(self, vf: _VF, control: _Control) -> Y:
 
         The interaction between the vector field and control; a PyTree of structure
         $T$.
-
-        !!! note
-
-            This function must be bilinear.
         """
         pass
 
@@ -258,6 +260,8 @@ def _prod(vf, control):
     return jnp.tensordot(vf, control, axes=jnp.ndim(control))
 
 
+# This class exists for backward compatibility with `WeaklyDiagonalControlTerm`. If we
+# were writing things again today it would be folded into just `ControlTerm`.
 class _AbstractControlTerm(AbstractTerm[_VF, _Control]):
     vector_field: Callable[[RealScalarLike, Y, Args], _VF]
     control: Union[
@@ -288,40 +292,143 @@ def to_ode(self) -> ODETerm:
 - `vector_field`: A callable representing the vector field. This callable takes three
     arguments `(t, y, args)`. `t` is a scalar representing the integration time. `y` is
     the evolving state of the system. `args` are any static arguments as passed to
-    [`diffrax.diffeqsolve`][].
-- `control`: The control. Should either be (A) a [`diffrax.AbstractPath`][], in which
-    case its `evaluate(t0, t1)` method will be used to give the increment of the control
-    over a time interval `[t0, t1]`, or (B) a callable `(t0, t1) -> increment`, which
-    returns the increment directly.
+    [`diffrax.diffeqsolve`][]. This `vector_field` can either be
+
+    1. a function that returns a PyTree of JAX arrays, or
+    2. it can return a
+        [Lineax linear operator](https://docs.kidger.site/lineax/api/operators),
+        as described above.
+
+- `control`: The control. Should either be
+
+    1. a [`diffrax.AbstractPath`][], in which case its `.evaluate(t0, t1)` method
+        will be used to give the increment of the control over a time interval
+        `[t0, t1]`, or
+    2. a callable `(t0, t1) -> increment`, which returns the increment directly.
 """
 
 
 class ControlTerm(_AbstractControlTerm[_VF, _Control]):
     r"""A term representing the general case of $f(t, y(t), args) \mathrm{d}x(t)$, in
-    which the vector field - control interaction is a matrix-vector product.
+    which the vector field ($f$) - control ($\mathrm{d}x$) interaction is a
+    matrix-vector product.
 
-    `vector_field` and `control` should both return PyTrees, both with the same
-    structure as the initial state `y0`. Every dimension of `control` is then
-    contracted against the last dimensions of `vector_field`; that is to say if each
-    leaf of `y0` has shape `(y1, ..., yN)`, and the corresponding leaf of `control`
-    has shape `(c1, ..., cM)`, then the corresponding leaf of `vector_field` should
-    have shape `(y1, ..., yN, c1, ..., cM)`.
+    This is typically used for either stochastic differential equations or for
+    controlled differential equations.
 
-    A common special case is when `y0` and `control` are vector-valued, and
-    `vector_field` is matrix-valued.
+    `ControlTerm` can be used in two different ways.
 
-    !!! example
+    1. Simple way: directly return JAX arrays.
+
+        `vector_field` and `control` should both return PyTrees, both with the same
+        structure as the initial state `y0`. All leaves should be JAX arrays.
+
+        If each leaf of `y0` has shape `(y1, ..., yN)`, and the corresponding leaf of
+        `control` has shape `(c1, ..., cM)`, then the corresponding leaf of
+        `vector_field` should have shape `(y1, ..., yN, c1, ..., cM)`. Leaf-by-leaf, the
+        corresponding dimensions of `vector_field` and control are contracted against
+        each other.
+
+        This includes normal matrix-vector products as a special case: when `y0` is an
+        array with shape `(m,)`, the control is an array with shape `(n,)`, and the
+        vector field is an array with shape `(m, n)`.
+
+    2. Advanced way: have the vector field return a [Lineax linear operator](https://docs.kidger.site/lineax/api/operators).
+
+        This is suitable for use cases in which you know that the vector field has
+        special structure -- e.g. it is diagonal -- and you would like to use that
+        structure for a more efficient implementation.
+
+        In this case, then `vector_field` should return a
+        [Lineax linear operator](https://docs.kidger.site/lineax/api/operators), the
+        control can return anything compatible with the
+        [`.mv`](https://docs.kidger.site/lineax/api/operators/#lineax.AbstractLinearOperator.mv)
+        method of that operator, and the interaction is defined as
+        `vector_field(t0, y, arg).mv(control(t0, t1))`.
+
+        In this case no special PyTree handling is done -- perform this inside the
+        operator's `.mv` if required. (As you can see, this approach is basically about
+        deferring the whole linear operation to Lineax.)
+
+    !!! Example
+
+        In this example we consider an SDE with `m`-dimensional state
+        $y \in \mathbb{R}^m$, an `n`-dimensional Brownian motion
+        $W(t) \in \mathbb{R}^n$, and a constant diffusion of shape `(m, n)`.
+
+        $\mathrm{d}y(t) = \begin{bmatrix} 1 & ... & 1 \\ & ... & \\ 1 & ... & 1 \end{bmatrix} \mathrm{d}W(t)$
 
         ```python
+        from diffrax import ControlTerm, diffeqsolve, UnsafeBrownianPath
+
+        y0 = jnp.ones((m,))
+        control = UnsafeBrownianPath(shape=(n,), key=...)
+
+        def vector_field(t, y, args):
+            return jnp.ones((m, n))
+
+        diffusion_term = ControlTerm(vector_field, control)
+        diffeqsolve(terms=diffusion_term, y0=y0, ...)
+        ```
+    !!! Example
+
+        In this example we consider an SDE with a one-dimensional state
+        $y(t) \in \mathbb{R}$ and a two-dimensional Brownian motion
+        $W(t) \in \mathbb{R}^2$, given by:
+
+        $\mathrm{d}y(t) = \begin{bmatrix} y(t) \\ y(t) + 1 \end{bmatrix} \mathrm{d}W(t)$
+
+        We use the simple matrix-vector product way of combining things.
+
+        ```python
+        from diffrax import ControlTerm, diffeqsolve, UnsafeBrownianPath
+
         control = UnsafeBrownianPath(shape=(2,), key=...)
-        vector_field = lambda t, y, args: jnp.stack([y, y], axis=-1)
+
+        def vector_field(t, y, args):
+            return jnp.stack([y, y + 1], axis=-1)
+
         diffusion_term = ControlTerm(vector_field, control)
         diffeqsolve(diffusion_term, ...)
         ```
 
-    !!! example
+    !!! Example
+
+        In this example we consider an SDE with two-dimensional state
+        $(y_1(t), y_2(t)) \in \mathbb{R}^2$ and a two-dimensional Brownian motion
+        $W(t) \in \mathbb{R}^2$ -- and for which the diffusion matrix is
+        diagonal.
+
+        $\mathrm{d}\begin{bmatrix} y_1 \\ y_2 \end{bmatrix}(t) = \begin{bmatrix} y_2(t) & 0 \\ 0 & y_1(t) \end{bmatrix} \mathrm{d}W(t)$
+
+        As such we use the more-advanced approach of using
+        [Lineax](https://github.com/patrick-kidger/lineax/)'s linear operators to
+        represent the diffusion matrix.
+
+        ```python
+        from diffrax import ControlTerm, diffeqsolve, UnsafeBrownianPath
+
+        control = UnsafeBrownianPath(shape=(2,), key=...)
+
+        def vector_field(t, y, args):
+            # y is a JAX array of shape (2,)
+            y1, y2 = y
+            diagonal = jnp.array([y2, y1])
+            return lineax.DiagonalLinearOperator(diagonal)
+
+        diffusion_term = ControlTerm(vector_field, control)
+        diffeqsolve(diffusion_term, ...)
+        ```
+
+    !!! Example
+
+        In this example we consider a controlled differnetial equation, for which the
+        control is given by an interpolation of some data. (See also the
+        [neural controlled differential equation](../examples/neural_cde/) example.)
 
         ```python
+        from diffrax import ControlTerm, diffeqsolve, LinearInterpolation, UnsafeBrownianPath
+
         ts = jnp.array([1., 2., 2.5, 3.])
         data = jnp.array([[0.1, 2.0],
                           [0.3, 1.5],
@@ -332,14 +439,29 @@ class ControlTerm(_AbstractControlTerm[_VF, _Control]):
         cde_term = ControlTerm(vector_field, control)
         diffeqsolve(cde_term, ...)
         ```
-    """
+    """  # noqa: E501
 
     def prod(self, vf: _VF, control: _Control) -> Y:
-        return jtu.tree_map(_prod, vf, control)
+        if isinstance(vf, lx.AbstractLinearOperator):
+            return vf.mv(control)
+        else:
+            return jtu.tree_map(_prod, vf, control)
 
 
 class WeaklyDiagonalControlTerm(_AbstractControlTerm[_VF, _Control]):
-    r"""A term representing the case of $f(t, y(t), args) \mathrm{d}x(t)$, in
+    r"""
+    DEPRECATED. Prefer:
+
+    ```python
+    def vector_field(t, y, args):
+        return lineax.DiagonalLinearOperator(...)
+
+    diffrax.ControlTerm(vector_field, ...)
+    ```
+
+    ---
+
+    A term representing the case of $f(t, y(t), args) \mathrm{d}x(t)$, in
     which the vector field - control interaction is a matrix-vector product, and the
     matrix is square and diagonal. In this case we may represent the matrix as a vector
     of just its diagonal elements. The matrix-vector product may be calculated by
@@ -360,6 +482,35 @@ class WeaklyDiagonalControlTerm(_AbstractControlTerm[_VF, _Control]):
         without the "weak". (This stronger property is useful in some SDE solvers.)
     """
 
+    def __check_init__(self):
+        warnings.warn(
+            "`WeaklyDiagonalControlTerm` is now deprecated, in favour combining "
+            "`ControlTerm` with a `lineax.AbstractLinearOperator`. This offers a way "
+            "to define a vector field with any kind of structure -- diagonal or "
+            "otherwise.\n"
+            "For a diagonal linear operator, then this can be easily converted as "
+            "follows. What was previously:\n"
+            "```\n"
+            "def vector_field(t, y, args):\n"
+            "    ...\n"
+            "    return some_vector\n"
+            "\n"
+            "diffrax.WeaklyDiagonalControlTerm(vector_field)\n"
+            "```\n"
+            "is now:\n"
+            "```\n"
+            "import lineax\n"
+            "\n"
+            "def vector_field(t, y, args):\n"
+            "    ...\n"
+            "    return lineax.DiagonalLinearOperator(some_vector)\n"
+            "\n"
+            "diffrax.ControlTerm(vector_field)\n"
+            "```\n"
+            "Lineax is available at `https://github.com/patrick-kidger/lineax`.\n",
+            stacklevel=3,
+        )
+
     def prod(self, vf: _VF, control: _Control) -> Y:
         with jax.numpy_dtype_promotion("standard"):
             return jtu.tree_map(operator.mul, vf, control)
diff --git a/docs/api/brownian.md b/docs/api/brownian.md
index c843f005..6d9719bf 100644
--- a/docs/api/brownian.md
+++ b/docs/api/brownian.md
@@ -22,3 +22,38 @@ SDEs are simulated using a Brownian motion as a control. (See the neural SDE exa
         members:
             - __init__
             - evaluate
+
+---
+
+## Lévy areas
+
+Brownian controls can return certain types of Lévy areas. These are iterated integrals
+of the Brownian motion, and are used by some SDE solvers. When a solver requires a 
+Lévy area, it will have a `minimal_levy_area` attribute, which will always return an
+abstract Lévy area type, and it can accept any subclass of that type.
+The inheritance hierarchy is as follows:
+```
+AbstractBrownianIncrement
+│   └── BrownianIncrement
+└── AbstractSpaceTimeLevyArea
+    │   └── SpaceTimeLevyArea
+    └── AbstractSpaceTimeTimeLevyArea
+            └── SpaceTimeTimeLevyArea
+```
+For example if `solver.minimal_levy_area` returns an `AbstractSpaceTimeLevyArea`, then
+the Brownian motion (which is either an `UnsafeBrownianPath` or 
+a `VirtualBrownianTree`) should be initialized with `levy_area=SpaceTimeLevyArea` or 
+`levy_area=SpaceTimeTimeLevyArea`. Note that for the Brownian motion,
+a concrete class must be used, not its abstract parent.
+
+::: diffrax.AbstractBrownianIncrement
+
+::: diffrax.BrownianIncrement
+
+::: diffrax.AbstractSpaceTimeLevyArea
+
+::: diffrax.SpaceTimeLevyArea
+
+::: diffrax.AbstractSpaceTimeTimeLevyArea
+
+::: diffrax.SpaceTimeTimeLevyArea
\ No newline at end of file
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/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
diff --git a/docs/api/solvers/sde_solvers.md b/docs/api/solvers/sde_solvers.md
index 66c09299..ecc9fbf9 100644
--- a/docs/api/solvers/sde_solvers.md
+++ b/docs/api/solvers/sde_solvers.md
@@ -6,7 +6,7 @@ See also [How to choose a solver](../../usage/how-to-choose-a-solver.md#stochast
 
     The type of solver chosen determines how the `terms` argument of `diffeqsolve` should be laid out.
     
-    Most solvers handle both ODEs and SDEs in the same way, and expect a single term. So for an ODE you would pass `terms=ODETerm(vector_field)`, and for an SDE you would pass `terms=MultiTerm(ODETerm(drift), ControlTerm(diffusion, brownian_motion))` or `terms=MultiTerm(ODETerm(drift), WeaklyDiagonalControlTerm(diffusion, brownian_motion))`. For example:
+    Most solvers handle both ODEs and SDEs in the same way, and expect a single term. So for an ODE you would pass `terms=ODETerm(vector_field)`, and for an SDE you would pass `terms=MultiTerm(ODETerm(drift), ControlTerm(diffusion, brownian_motion))`. For example:
 
     ```python
     drift = lambda t, y, args: -y
@@ -18,7 +18,7 @@ See also [How to choose a solver](../../usage/how-to-choose-a-solver.md#stochast
 
     For any individual solver then this is documented below, and is also available programatically under `<solver>.term_structure`.
 
-    For advanced users, note that we typically accept any `AbstractTerm` for the diffusion, so it could be a custom one that implements more-efficient behaviour for the structure of your diffusion matrix. (Much like how [`diffrax.WeaklyDiagonalControlTerm`][] is more efficient than [`diffrax.ControlTerm`][] for diagonal diffusions.)
+    For advanced users, note that we typically accept any `AbstractTerm` for the diffusion, so it could be a custom one that implements more-efficient behaviour for the structure of your diffusion matrix.
 
 ---
 
@@ -52,7 +52,7 @@ These solvers can be used to solve SDEs just as well as they can be used to solv
 
 !!! info "Term structure"
 
-    These solvers are SDE-specific. For these, `terms` must specifically be of the form `MultiTerm(ODETerm(...), SomeOtherTerm(...))` (Typically `SomeOTherTerm` will be a `ControlTerm` or `WeaklyDiagonalControlTerm`) representing the drift and diffusion specifically.
+    These solvers are SDE-specific. For these, `terms` must specifically be of the form `MultiTerm(ODETerm(...), SomeOtherTerm(...))` (Typically `SomeOTherTerm` will be a `ControlTerm` representing the drift and diffusion specifically.
 
 
 ::: diffrax.EulerHeun
diff --git a/docs/api/terms.md b/docs/api/terms.md
index b78e5241..553b3fbb 100644
--- a/docs/api/terms.md
+++ b/docs/api/terms.md
@@ -4,7 +4,7 @@ One of the advanced features of Diffrax is its *term* system. When we write down
 
 $\mathrm{d}y(t) = f(t, y(t))\mathrm{d}t + g(t, y(t))\mathrm{d}w(t)$
 
-then we have two "terms": a drift and a diffusion. Each of these terms has two parts: a *vector field* ($f$ or $g$) and a *control* ($\mathrm{d}t$ or $\mathrm{d}w(t)$). There is also an implicit assumption about how the vector field and control interact: $f$ and $\mathrm{d}t$ interact as a vector-scalar product. $g$ and $\mathrm{d}w(t)$ interact as a matrix-vector product. (This interaction is always linear.)
+then we have two "terms": a drift and a diffusion. Each of these terms has two parts: a *vector field* ($f$ or $g$) and a *control* ($\mathrm{d}t$ or $\mathrm{d}w(t)$). In addition (often not represented in mathematical notation), there is also a choice of how the vector field and control interact: $f$ and $\mathrm{d}t$ interact as a vector-scalar product. $g$ and $\mathrm{d}w(t)$ interact as a matrix-vector product. (In general this interaction is always bilinear.)
 
 "Terms" are thus the building blocks of differential equations.
 
@@ -12,13 +12,17 @@ then we have two "terms": a drift and a diffusion. Each of these terms has two p
 
     Consider the ODE $\frac{\mathrm{d}{y}}{\mathrm{d}t} = f(t, y(t))$. Then this has vector field $f$, control $\mathrm{d}t$, and their interaction is a vector-scalar product. This can be described as a single [`diffrax.ODETerm`][].
 
-If multiple terms affect the same evolving state, then they should be grouped into a single [`diffrax.MultiTerm`][].
+#### Adding multiple terms, such as SDEs
+
+We can add multiple terms together by grouping them into a single [`diffrax.MultiTerm`][].
 
 !!! example
 
-    An SDE would have its drift described by [`diffrax.ODETerm`][] and the diffusion described by a [`diffrax.ControlTerm`][]. As these affect the same evolving state variable, they should be passed to the solver as `MultiTerm(ODETerm(...), ControlTerm(...))`.
+    The SDE above would have its drift described by [`diffrax.ODETerm`][] and the diffusion described by a [`diffrax.ControlTerm`][]. As these affect the same evolving state variable, they should be passed to the solver as `MultiTerm(ODETerm(...), ControlTerm(...))`.
+
+#### Independent terms, such as Hamiltonian systems
 
-If terms affect different pieces of the state, then they should be placed in some PyTree structure. (The exact structure will depend on what the solver accepts.)
+If terms affect different pieces of the state, then they should be placed in some PyTree structure.
 
 !!! example
 
@@ -28,7 +32,31 @@ If terms affect different pieces of the state, then they should be placed in som
 
     These would be passed to the solver as the 2-tuple of `(ODETerm(...), ODETerm(...))`.
 
-Each solver is capable of handling certain classes of problems, as described by their `solver.term_structure`.
+#### What each solver accepts
+
+Each solver in Diffrax will specify what kinds of problems it can handle, as described by their `.term_structure` attribute. Not all solvers are able to handle all problems!
+
+Some example term structures include:
+
+1. `solver.term_structure = AbstractTerm`
+
+    In this case the solver can handle a simple ODE as descibed above: `ODETerm` is a subclass of `AbstractTerm`.
+
+    It can also handle SDEs: `MultiTerm(ODETerm(...), ControlTerm(...))` includes everything wrapped into a single term (the `MultiTerm`), and at that point this defines an interface the solver knows how to handle.
+
+    Most solvers in Diffrax have this term structure.
+
+2. `solver.term_structure = MultiTerm[tuple[ODETerm, ControlTerm]]`
+
+    In this case the solver specifically handles just SDEs of the form `MultiTerm(ODETerm(...), ControlTerm(...))`; nothing else is compatible.
+
+    Some SDE-specific solvers have this term structure.
+
+3. `solver.term_structure = (AbstractTerm, AbstractTerm)`
+
+    In this case the solver is used to solve ODEs like the Hamiltonian system described above: we have a PyTree of terms, each of which is treated individually.
+
+---
 
 ??? abstract "`diffrax.AbstractTerm`"
 
@@ -41,10 +69,12 @@ Each solver is capable of handling certain classes of problems, as described by
                 - vf_prod
                 - is_vf_expensive
 
----
+??? note "Defining your own term types"
+
+    For advanced users: you can create your own terms if appropriate. For example if your diffusion is matrix, itself computed as a matrix-matrix product, then you may wish to define a custom term and specify its [`diffrax.AbstractTerm.vf_prod`][] method. By overriding this method you could express the contraction of the vector field - control as a matrix-(matix-vector) product, which is more efficient than the default (matrix-matrix)-vector product.
+
 
-!!! note
-    You can create your own terms if appropriate: e.g. if a diffusion matrix has some particular structure, and you want to use a specialised more efficient matrix-vector product algorithm in `prod`. For example this is what [`diffrax.WeaklyDiagonalControlTerm`][] does, as compared to just [`diffrax.ControlTerm`][].
+---
 
 ::: diffrax.ODETerm
     selection:
@@ -57,12 +87,6 @@ Each solver is capable of handling certain classes of problems, as described by
             - __init__
             - to_ode
 
-::: diffrax.WeaklyDiagonalControlTerm
-    selection:
-        members:
-            - __init__
-            - to_ode
-
 ::: diffrax.MultiTerm
     selection:
         members:
diff --git a/docs/devdocs/srk_example.ipynb b/docs/devdocs/srk_example.ipynb
index f1905822..aa2a38cf 100644
--- a/docs/devdocs/srk_example.ipynb
+++ b/docs/devdocs/srk_example.ipynb
@@ -19,11 +19,11 @@
     "\n",
     "To account for time-dependent noise, the SRK adds a term to the output of each step, which allows it to still maintain its usual strong order of convergence.\n",
     "\n",
-    "The SRK is capable of utilising various types of time Levy area, depending on the tableau provided. It can use:\n",
+    "The SRK is capable of utilising various types of time Lévy area, depending on the tableau provided. It can use:\n",
     "\n",
-    "- just the Brownian motion $W$, without any Levy area\n",
-    "- $W$ and the space-time Levy area $H$\n",
-    "- $W$, $H$ and the space-time-time Levy area $K$.\n",
+    "- just the Brownian motion $W$, without any Lévy area\n",
+    "- $W$ and the space-time Lévy area $H$\n",
+    "- $W$, $H$ and the space-time-time Lévy area $K$.\n",
     "For more information see the documentation of the `StochasticButcherTableau` class.\n",
     "\n",
     "First we will demonstrate an additive-noise-only SRK method, the ShARK method, on an SDE with additive, time-dependent noise.\n",
@@ -263,7 +263,7 @@
     "    \n",
     "\n",
     "### Shifted Additive-noise Euler (SEA)\n",
-    "This variant of the Euler-Maruyama makes use of the space-time Levy area, which improves its local error to $O(h^2)$ compared to $O(h^{1.5})$ of the standard Euler-Maruyama. Nevertheless, it has a strong order of only 1 for additive-noise SDEs.\n",
+    "This variant of the Euler-Maruyama makes use of the space-time Lévy area, which improves its local error to $O(h^2)$ compared to $O(h^{1.5})$ of the standard Euler-Maruyama. Nevertheless, it has a strong order of only 1 for additive-noise SDEs.\n",
     "\n",
     "\n",
     " ### The \"Splitting Path Runge-Kutta\" (SPaRK) method\n",
@@ -277,7 +277,7 @@
     "When the noise is commutative it has order 1.\n",
     "When the noise is additive it has order 1.5.\n",
     "For the Langevin SDE it has order 2.\n",
-    "Requires the space-time Levy area H.\n",
+    "Requires the space-time Lévy area H.\n",
     "It also natively supports adaptive time-stepping.\n",
     "\n",
     "\n",
diff --git a/docs/usage/extending.md b/docs/usage/extending.md
index 12648445..9cf783f6 100644
--- a/docs/usage/extending.md
+++ b/docs/usage/extending.md
@@ -24,7 +24,7 @@ The main points of extension are as follows:
 - **Custom controls** (e.g. **custom interpolation schemes** analogous to [`diffrax.CubicInterpolation`][]) should inherit from [`diffrax.AbstractPath`][].
 
 - **Custom terms** should inherit from [`diffrax.AbstractTerm`][].
-    - For example, if the vector field - control interaction is a matrix-vector product, but the matrix is known to have special structure, then you may wish to create a custom term that can calculate this interaction more efficiently than is given by a full matrix-vector product. For example this is done with [`diffrax.WeaklyDiagonalControlTerm`][] as compared to [`diffrax.ControlTerm`][].
+    - For example, if the vector field - control interaction is a matrix-vector product, but the matrix is known to have special structure, then you may wish to create a custom term that can calculate this interaction more efficiently than is given by a full matrix-vector product. Given the large suite of linear operators [lineax](https://docs.kidger.site/lineax/) implements (which are fully supported by [`diffrax.ControlTerm`][]), this is likely rarely necessary.
 
 In each case we recommend looking up existing solvers/etc. in Diffrax to understand how to implement them.
 
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,
diff --git a/pyproject.toml b/pyproject.toml
index 6444e742..434d6fbe 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "diffrax"
-version = "0.5.1"
+version = "0.6.0"
 description = "GPU+autodiff-capable ODE/SDE/CDE solvers written in JAX."
 readme = "README.md"
 requires-python ="~=3.9"
@@ -8,7 +8,7 @@ license = {file = "LICENSE"}
 authors = [
   {name = "Patrick Kidger", email = "contact@kidger.site"},
 ]
-keywords = ["jax", "dynamical-systems", "differential-equations", "deep-learning", "equinox", "neural-differential-equations"]
+keywords = ["jax", "dynamical-systems", "differential-equations", "deep-learning", "equinox", "neural-differential-equations", "diffrax"]
 classifiers = [
     "Development Status :: 3 - Alpha",
     "Intended Audience :: Developers",


Test patch:
diff --git a/test/helpers.py b/test/helpers.py
index b7e33b54..fd4097dd 100644
--- a/test/helpers.py
+++ b/test/helpers.py
@@ -159,7 +159,7 @@ def _batch_sde_solve(
         shape=struct,
         tol=bm_tol,
         key=key,
-        levy_area=concrete_la,  # pyright: ignore
+        levy_area=concrete_la,
     )
     terms = get_terms(bm)
     if controller is None:
@@ -183,12 +183,12 @@ def _resulting_levy_area(
     levy_area1: type[diffrax.AbstractBrownianIncrement],
     levy_area2: type[diffrax.AbstractBrownianIncrement],
 ) -> type[diffrax.AbstractBrownianIncrement]:
-    """A helper that returns the stricter Levy area.
+    """A helper that returns the stricter Lévy area.
 
     **Arguments:**
 
-    - `levy_area1`: The first Levy area type.
-    - `levy_area2`: The second Levy area type.
+    - `levy_area1`: The first Lévy area type.
+    - `levy_area2`: The second Lévy area type.
 
     **Returns:**
 
diff --git a/test/test_adjoint.py b/test/test_adjoint.py
index f61a3f5e..3d143593 100644
--- a/test/test_adjoint.py
+++ b/test/test_adjoint.py
@@ -8,6 +8,7 @@
 import jax.numpy as jnp
 import jax.random as jr
 import jax.tree_util as jtu
+import lineax as lx
 import optax
 import pytest
 from jaxtyping import Array
@@ -264,6 +265,7 @@ class ExponentialDecayToSteadyState(eqx.Module):
         non_jax_type: Any
 
         def __call__(self, t, y, args):
+            del t, args
             return self.steady_state - y
 
     def loss(model, target_steady_state):
@@ -275,7 +277,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,
@@ -286,7 +289,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)
@@ -329,7 +332,11 @@ def run(model):
     run(mlp)
 
 
-def test_sde_against(getkey):
+@pytest.mark.parametrize(
+    "diffusion_fn",
+    ["weak", "lineax"],
+)
+def test_sde_against(diffusion_fn, getkey):
     def f(t, y, args):
         k0, _ = args
         return -k0 * y
@@ -338,6 +345,10 @@ def g(t, y, args):
         _, k1 = args
         return k1 * y
 
+    def g_lx(t, y, args):
+        _, k1 = args
+        return lx.DiagonalLinearOperator(k1 * y)
+
     t0 = 0
     t1 = 1
     dt0 = 0.001
@@ -345,7 +356,10 @@ def g(t, y, args):
     shape = (2,)
     bm = diffrax.VirtualBrownianTree(t0, t1, tol, shape, key=getkey())
     drift = diffrax.ODETerm(f)
-    diffusion = diffrax.WeaklyDiagonalControlTerm(g, bm)
+    if diffusion_fn == "weak":
+        diffusion = diffrax.WeaklyDiagonalControlTerm(g, bm)
+    else:
+        diffusion = diffrax.ControlTerm(g_lx, bm)
     terms = diffrax.MultiTerm(drift, diffusion)
     solver = diffrax.Heun()
 
diff --git a/test/test_brownian.py b/test/test_brownian.py
index 70eeb91b..3a265019 100644
--- a/test/test_brownian.py
+++ b/test/test_brownian.py
@@ -13,6 +13,11 @@
 import scipy.stats as stats
 
 
+_levy_areas = (
+    diffrax.BrownianIncrement,
+    diffrax.SpaceTimeLevyArea,
+    diffrax.SpaceTimeTimeLevyArea,
+)
 _Spline: TypeAlias = Literal["quad", "sqrt", "zero"]
 _vals = {
     int: [0, 2],
@@ -31,56 +36,43 @@ def _make_struct(shape, dtype):
 @pytest.mark.parametrize(
     "ctr", [diffrax.UnsafeBrownianPath, diffrax.VirtualBrownianTree]
 )
-@pytest.mark.parametrize(
-    "levy_area", [diffrax.BrownianIncrement, diffrax.SpaceTimeLevyArea]
-)
+@pytest.mark.parametrize("levy_area", _levy_areas)
 @pytest.mark.parametrize("use_levy", (False, True))
 def test_shape_and_dtype(ctr, levy_area, use_levy, getkey):
     t0 = 0.0
     t1 = 2.0
 
-    shapes = (
-        (),
-        (0,),
-        (1, 0),
-        (2,),
-        (3, 4),
-        (1, 2, 3, 4),
-        (1, 2, 3, 4),
-        {
-            "a": (1,),
-            "b": (2, 3),
-        },
-        {
-            "a": (1,),
-            "b": (2, 3),
-        },
-        (
-            (1, 2),
-            (
-                (3, 4),
-                (5, 6),
-            ),
-        ),
+    shapes_dtypes1 = (
+        ((), None),
+        ((0,), None),
+        ((1, 0), None),
+        ((3, 4), jnp.float32),
+        ((1, 2, 3, 4), jnp.float64),
+        ({"a": (1,), "b": (2, 3)}, {"a": None, "b": jnp.float64}),
+        ((2,), jnp.float16),
+        (((1, 2), ((3, 4), (5, 6))), (jnp.float16, (jnp.float32, jnp.float64))),
     )
 
-    dtypes = (
-        None,
-        None,
-        None,
-        jnp.float16,
-        jnp.float32,
-        jnp.float64,
-        jnp.complex128,
-        {"a": None, "b": jnp.float64},
-        {"a": jnp.float64, "b": jnp.complex128},
-        (jnp.float16, (jnp.float32, jnp.float64)),
+    shapes_dtypes2 = (
+        ((1, 2, 3, 4), jnp.complex128),
+        ({"a": (1,), "b": (2, 3)}, {"a": jnp.float64, "b": jnp.complex128}),
     )
 
+    if (
+        ctr is diffrax.VirtualBrownianTree
+        and levy_area is diffrax.SpaceTimeTimeLevyArea
+    ):
+        # VBT with STTLA does not support complex dtypes
+        # because it uses jax.random.multivariate_normal
+        shapes_dtypes = shapes_dtypes1
+
+    else:
+        shapes_dtypes = shapes_dtypes1 + shapes_dtypes2
+
     def is_tuple_of_ints(obj):
         return isinstance(obj, tuple) and all(isinstance(x, int) for x in obj)
 
-    for shape, dtype in zip(shapes, dtypes):
+    for shape, dtype in shapes_dtypes:
         # Shape to pass as input
         if dtype is not None:
             shape = jtu.tree_map(_make_struct, shape, dtype, is_leaf=is_tuple_of_ints)
@@ -115,22 +107,26 @@ def is_tuple_of_ints(obj):
                     if isinstance(out, diffrax.SpaceTimeLevyArea):
                         h = out.H
                         assert eqx.filter_eval_shape(lambda: h) == shape
+                        if isinstance(out, diffrax.SpaceTimeTimeLevyArea):
+                            k = out.K
+                            assert eqx.filter_eval_shape(lambda: k) == shape
                 else:
                     w = out
                 assert eqx.filter_eval_shape(lambda: w) == shape
 
 
 @pytest.mark.parametrize(
-    "ctr", [diffrax.UnsafeBrownianPath, diffrax.VirtualBrownianTree]
-)
-@pytest.mark.parametrize(
-    "levy_area", [diffrax.BrownianIncrement, diffrax.SpaceTimeLevyArea]
+    "ctr", [diffrax.VirtualBrownianTree, diffrax.UnsafeBrownianPath]
 )
-@pytest.mark.parametrize("use_levy", (False, True))
+@pytest.mark.parametrize("levy_area", _levy_areas)
+@pytest.mark.parametrize("use_levy", (True, False))
 def test_statistics(ctr, levy_area, use_levy):
     # Deterministic key for this test; not using getkey()
     key = jr.PRNGKey(5678)
-    keys = jr.split(key, 10000)
+    num_samples = 60000
+    keys = jr.split(key, num_samples)
+    t0, t1 = 0.0, 5.0
+    dt = t1 - t0
 
     def _eval(key):
         if ctr is diffrax.UnsafeBrownianPath:
@@ -139,7 +135,7 @@ def _eval(key):
             path = ctr(t0=0, t1=5, tol=2**-5, shape=(), key=key, levy_area=levy_area)
         else:
             assert False
-        return path.evaluate(0, 5, use_levy=use_levy)
+        return path.evaluate(t0, t1, use_levy=use_levy)
 
     values = jax.vmap(_eval)(keys)
     if use_levy:
@@ -149,27 +145,149 @@ def _eval(key):
         if isinstance(values, diffrax.SpaceTimeLevyArea):
             h = values.H
             assert h is not None
-            assert h.shape == (10000,)
-            ref_dist = stats.norm(loc=0, scale=math.sqrt(5 / 12))
-            _, pval = stats.kstest(h, ref_dist.cdf)
+            assert h.shape == (num_samples,)
+            ref_dist_hh = stats.norm(loc=0, scale=math.sqrt(dt / 12))
+            _, pval = stats.kstest(h, ref_dist_hh.cdf)
+            assert pval > 0.1
+            # Check independence of h and w
+            assert jnp.mean(h * w) < 0.01
+
+        elif isinstance(values, diffrax.SpaceTimeTimeLevyArea):
+            h = values.H
+            k = values.K
+            assert h is not None
+            assert k is not None
+            assert h.shape == (num_samples,)
+            assert k.shape == (num_samples,)
+            ref_dist_hh = stats.norm(loc=0, scale=math.sqrt(dt / 12))
+            _, pval = stats.kstest(h, ref_dist_hh.cdf)
+            assert pval > 0.1
+
+            ref_dist_kk = stats.norm(loc=0, scale=math.sqrt(dt / 720))
+            _, pval = stats.kstest(k, ref_dist_kk.cdf)
             assert pval > 0.1
+
+            # Check independence of w, h and k
+            assert jnp.abs(jnp.mean(w * h)) < 0.01
+            assert jnp.abs(jnp.mean(h * k)) < 0.01
+            assert jnp.abs(jnp.mean(k * w)) < 0.01
+
     else:
         w = values
-    assert w.shape == (10000,)
-    ref_dist = stats.norm(loc=0, scale=math.sqrt(5))
+    assert w.shape == (num_samples,)
+    ref_dist = stats.norm(loc=0, scale=math.sqrt(dt))
     _, pval = stats.kstest(w, ref_dist.cdf)
     assert pval > 0.1
 
 
-def conditional_statistics(
+def _true_cond_stats_wh(bm_s, bm_u, s, r, u):
+    w_s = bm_s.W
+    w_u = bm_u.W
+    h_s = bm_s.H
+    h_u = bm_u.H
+    su = u - s
+    sr = r - s
+    ru = u - r
+    sr3 = jnp.power(sr, 3)
+    ru3 = jnp.power(ru, 3)
+    su2 = jnp.power(su, 2)
+    sr_ru_half = jnp.sqrt(sr * ru)
+    d = math.sqrt(sr3 + ru3)
+    a = (1 / (2 * su * d)) * sr3 * sr_ru_half
+    b = (1 / (2 * su * d)) * ru3 * sr_ru_half
+    c = (1.0 / (math.sqrt(12) * d)) * sr ** (3 / 2) * ru ** (3 / 2)
+    h_su = (1.0 / su) * (u * h_u - s * h_s - u / 2 * w_s + s / 2 * w_u)
+
+    # Check w_r|(w_s, w_u, h_s, h_u)
+    w_mean = w_s + (sr / su) * (w_u - w_s) + (6 * sr * ru / (su2)) * h_su
+    w_std = 2 * (a + b) / su
+
+    # Check h_r|(w_s, w_u, h_s, h_u)
+    h_mean = (s / r) * h_s + (sr3 / (r * su2)) * h_su + 0.5 * w_s - s / (2 * r) * w_mean
+    h_var = (c / r) ** 2 + ((a * u + s * b) / (r * su)) ** 2
+
+    h_std = jnp.sqrt(h_var)
+
+    return w_mean, w_std, h_mean, h_std
+
+
+def _true_cond_stats_whk(bm_s, bm_u, s, r, u):
+    su = u - s
+    sr = r - s
+    ru = u - r
+    sr3 = jnp.power(sr, 3)
+    su3 = jnp.power(su, 3)
+    w_s = bm_s.W
+    w_u = bm_u.W
+    h_s = bm_s.H
+    h_u = bm_u.H
+    k_s = bm_s.K
+    k_u = bm_u.K
+
+    su5 = jnp.power(su, 5)
+    sr2 = jnp.square(sr)
+    ru2 = jnp.square(ru)
+    bh_s = s * h_s
+    bh_u = u * h_u
+    bk_s = s**2 * k_s
+    bk_u = u**2 * k_u
+    # u_bb_s := u * brownian bridge on [0,u] evaluated at s
+    u_bb_s = u * w_s - s * w_u
+
+    # Chen's relation for H
+    bh_su = bh_u - bh_s - 0.5 * u_bb_s
+
+    # Chen's relation for \bar{K}_{s,u} := (u-s)^2 * K_{s,u}
+    bk_su = bk_u - bk_s - su / 2 * bh_s + s / 2 * bh_su - ((u - 2 * s) / 12) * u_bb_s
+
+    # compute the mean of (W_sr, H_sr, K_sr) conditioned on
+    # (W_s, H_s, K_s, W_u, H_u, K_u)
+    bb_mean = (6 * sr * ru / su3) * bh_su + (
+        120 * sr * ru * (su / 2 - sr) / su5
+    ) * bk_su
+    mean_w = (sr / su) * (w_u - w_s) + bb_mean
+    mean_hh = (sr2 / su3) * bh_su + (30 * sr2 * ru / su5) * bk_su
+    mean_kk = (sr3 / su5) * bk_su
+
+    mean_whk = jnp.stack([mean_w, mean_hh, mean_kk], axis=0)
+
+    # now compute the covariance matrix of (W_sr, H_sr, K_sr) conditioned on
+    # (W_s, H_s, K_s, W_u, H_u, K_u)
+    # note that the covariance matrix is independent of the values of
+    # (W_s, H_s, K_s, W_u, H_u, K_u), since those are already represented
+    # in the mean.
+    sr5 = jnp.power(sr, 5)
+
+    ww_cov = (sr * ru * ((sr - ru) ** 4 + 4 * (sr2 * ru2))) / su5
+    wh_cov = -(sr3 * ru * (sr2 - 3 * sr * ru + 6 * ru2)) / (2 * su5)
+    wk_cov = (sr**4 * ru * (sr - ru)) / (12 * su5)
+    hh_cov = sr / 12 * (1 - (sr3 * (sr2 + 2 * sr * ru + 16 * ru2)) / su5)
+    hk_cov = -(sr5 * ru) / (24 * su5)
+    kk_cov = sr / 720 * (1 - sr5 / su5)
+
+    cov = jnp.array(
+        [
+            [ww_cov, wh_cov, wk_cov],
+            [wh_cov, hh_cov, hk_cov],
+            [wk_cov, hk_cov, kk_cov],
+        ]
+    )
+
+    return mean_whk, cov
+
+
+def _conditional_statistics(
     levy_area, use_levy: bool, tol, spacing, spline: _Spline, min_num_points
 ):
     key = jr.PRNGKey(5678)
     bm_key, sample_key, permute_key = jr.split(key, 3)
     # Get some randomly selected points; not too close to avoid discretisation error.
-    t0 = 0.3
+    t0 = 0.0
     t1 = 8.7
-    ts = jr.uniform(sample_key, shape=(100,), minval=t0, maxval=t1)
+    boundary = 0.1
+    ts = jr.uniform(
+        sample_key, shape=(100,), minval=t0 + boundary, maxval=t1 - boundary
+    )
     sorted_ts = jnp.sort(ts)
     ts = []
     prev_ti = sorted_ts[0]
@@ -184,14 +302,14 @@ def conditional_statistics(
     ts = jr.permutation(permute_key, ts)
 
     # Get some random paths
-    bm_keys = jr.split(bm_key, 10000)
+    num_paths = 20000
+    bm_keys = jr.split(bm_key, num_paths)
 
     path = jax.vmap(
         lambda k: diffrax.VirtualBrownianTree(
             t0=t0, t1=t1, shape=(), tol=tol, key=k, levy_area=levy_area, _spline=spline
         )
     )(bm_keys)
-
     # Sample some points
     out = []
     for ti in ts:
@@ -202,6 +320,10 @@ def conditional_statistics(
     pvals_w1 = []
     pvals_w2 = []
     pvals_h = []
+    pvals_k = []
+
+    total_mean_err = jnp.zeros((3,), dtype=jnp.float64)
+    total_cov_err = jnp.zeros((3, 3), dtype=jnp.float64)
 
     # Test their conditional statistics
     for i in range(1, len(ts) - 1):
@@ -212,21 +334,26 @@ def conditional_statistics(
             w_s = bm_s.W
             w_r = bm_r.W
             w_u = bm_u.W
-            if levy_area is diffrax.SpaceTimeLevyArea:
+            if issubclass(levy_area, diffrax.AbstractSpaceTimeLevyArea):
                 h_s = bm_s.H
                 h_r = bm_r.H
-                h_u = bm_u.H
             else:
                 h_s = None
                 h_r = None
-                h_u = None
+            if issubclass(levy_area, diffrax.AbstractSpaceTimeTimeLevyArea):
+                k_s = bm_s.K
+                k_r = bm_r.K
+            else:
+                k_s = None
+                k_r = None
         else:
             w_s = bm_s
             w_r = bm_r
             w_u = bm_u
             h_s = None
             h_r = None
-            h_u = None
+            k_s = None
+            k_r = None
 
         # Check w_r|(w_s, w_u)
         w_mean1 = w_s + (w_u - w_s) * ((r - s) / (u - s))
@@ -238,50 +365,102 @@ def conditional_statistics(
         # multiple-testing correction.
         pvals_w1.append(pval_w1)
 
-        if levy_area is diffrax.SpaceTimeLevyArea and use_levy:
-            assert h_s is not None
-            assert h_r is not None
-            assert h_u is not None
-            s = s - t0
-            r = r - t0
-            u = u - t0
-            su = u - s
-            sr = r - s
-            ru = u - r
-            d = math.sqrt(sr**3 + ru**3)
-            a = (1 / (2 * su * d)) * sr ** (7 / 2) * math.sqrt(ru)
-            b = (1 / (2 * su * d)) * ru ** (7 / 2) * math.sqrt(sr)
-            c = (1.0 / (math.sqrt(12) * d)) * sr ** (3 / 2) * ru ** (3 / 2)
-            h_su = (1.0 / su) * (u * h_u - s * h_s - u / 2 * w_s + s / 2 * w_u)
+        s = s - t0
+        r = r - t0
+        u = u - t0
+        sr = r - s
+        if levy_area == diffrax.SpaceTimeTimeLevyArea and use_levy:
+            assert bm_s.H is not None
+            assert bm_r.H is not None
+            assert bm_u.H is not None
+            assert bm_s.K is not None
+            assert bm_r.K is not None
+            assert bm_u.K is not None
+            sr2 = jnp.square(sr)
+            bh_s = s * h_s
+            bh_r = r * h_r
+            bk_s = s**2 * k_s
+            bk_r = r**2 * k_r
+
+            # Compute the target conditional mean and covariance
+            true_mean_whk, true_cov = _true_cond_stats_whk(bm_s, bm_u, s, r, u)
+
+            # now compute the values of (W_sr, H_sr, K_sr), which are to be tested
+            # against the normal distribution N(mean, cov)
+            w_sr = w_r - w_s
+            r_bb_s = r * w_s - s * w_r
+            bh_sr = bh_r - bh_s - 0.5 * r_bb_s
+            h_sr = bh_sr / sr
+            bk_sr = (
+                bk_r
+                - bk_s
+                - 0.5 * (sr * bh_s - s * bh_sr)
+                - ((r - 2 * s) / 12) * r_bb_s
+            )
+            k_sr = bk_sr / sr2
+
+            y = jnp.stack([w_sr, h_sr, k_sr], axis=0)
+
+            # now we have to confirm that (w_centred, h_centred, k_centred) have
+            # zero mean and covariance matrix cov
+
+            hat_y = y - true_mean_whk
+            tilde_mean = jnp.mean(true_mean_whk, axis=1)
+            y_mean = jnp.mean(y, axis=1)
+            mean_diff = y_mean - tilde_mean
+            emp_cov = jnp.cov(hat_y)
+
+            mean_err = jnp.abs(mean_diff)
+            cov_err = jnp.abs(emp_cov - true_cov)
+            total_mean_err += mean_err / (len(ts) - 2)
+            total_cov_err += cov_err / (len(ts) - 2)
+
+            hat_w, hat_h, hat_k = hat_y
+            w_var2, h_var2, k_var2 = jnp.diag(true_cov)
+
+            normalised_w2 = hat_w / math.sqrt(w_var2)
+            _, pval_w2 = stats.kstest(normalised_w2, stats.norm.cdf)
+            pvals_w2.append(pval_w2)
+
+            normalised_h = hat_h / math.sqrt(h_var2)
+            _, pval_h = stats.kstest(normalised_h, stats.norm.cdf)
+            pvals_h.append(pval_h)
+
+            normalised_k = hat_k / math.sqrt(k_var2)
+            _, pval_k = stats.kstest(normalised_k, stats.norm.cdf)
+            pvals_k.append(pval_k)
+
+        elif levy_area == diffrax.SpaceTimeLevyArea and use_levy:
+            assert bm_s.H is not None
+            assert bm_r.H is not None
+            assert bm_u.H is not None
+
+            # Compute the true conditional statistics for W and H
+            w_mean2, w_std2, h_mean, h_std = _true_cond_stats_wh(bm_s, bm_u, s, r, u)
 
             # Check w_r|(w_s, w_u, h_s, h_u)
-            w_mean2 = w_s + (sr / su) * (w_u - w_s) + (6 * sr * ru / (su**2)) * h_su
-            w_std2 = 2 * (a + b) / su
             normalised_w2 = (w_r - w_mean2) / w_std2
             _, pval_w2 = stats.kstest(normalised_w2, stats.norm.cdf)
             pvals_w2.append(pval_w2)
 
             # Check h_r|(w_s, w_u, h_s, h_u)
-            h_mean = (
-                (s / r) * h_s
-                + (sr**3 / (r * su**2)) * h_su
-                + 0.5 * w_s
-                - s / (2 * r) * w_mean2
-            )
-            h_var = (c / r) ** 2 + ((a * u + s * b) / (r * su)) ** 2
-            h_std = math.sqrt(h_var)
             normalised_hh = (h_r - h_mean) / h_std
             _, pval_h = stats.kstest(normalised_hh, stats.norm.cdf)
             pvals_h.append(pval_h)
-    return jnp.array(pvals_w1), jnp.array(pvals_w2), jnp.array(pvals_h)
+    return (
+        jnp.array(pvals_w1),
+        jnp.array(pvals_w2),
+        jnp.array(pvals_h),
+        jnp.array(pvals_k),
+        total_mean_err,
+        total_cov_err,
+    )
 
 
-@pytest.mark.parametrize(
-    "levy_area", [diffrax.BrownianIncrement, diffrax.SpaceTimeLevyArea]
-)
-@pytest.mark.parametrize("use_levy", (False, True))
+@pytest.mark.parametrize("levy_area", _levy_areas)
+@pytest.mark.parametrize("use_levy", (True, False))
 def test_conditional_statistics(levy_area, use_levy):
-    pvals_w1, pvals_w2, pvals_h = conditional_statistics(
+    pvals_w1, pvals_w2, pvals_h, pvals_k, mean_err, cov_err = _conditional_statistics(
         levy_area,
         use_levy,
         tol=2**-10,
@@ -290,32 +469,55 @@ def test_conditional_statistics(levy_area, use_levy):
         min_num_points=90,
     )
     assert jnp.all(pvals_w1 > 0.1 / pvals_w1.shape[0])
-    if levy_area is diffrax.SpaceTimeLevyArea and use_levy:
+
+    if levy_area is diffrax.SpaceTimeTimeLevyArea and use_levy:
+        assert jnp.sum(mean_err) < 0.005
+        assert jnp.sum(cov_err) < 0.001
         assert jnp.all(pvals_w2 > 0.1 / pvals_w2.shape[0])
         assert jnp.all(pvals_h > 0.1 / pvals_h.shape[0])
+        assert jnp.all(pvals_k > 0.1 / pvals_k.shape[0])
+
+    elif levy_area is diffrax.SpaceTimeLevyArea and use_levy:
+        assert jnp.all(pvals_w2 > 0.1 / pvals_w2.shape[0])
+        assert jnp.all(pvals_h > 0.1 / pvals_h.shape[0])
+        assert len(pvals_k) == 0
+        assert jnp.all(mean_err == 0.0)
+        assert jnp.all(cov_err == 0.0)
+
     else:
         assert len(pvals_w2) == 0
         assert len(pvals_h) == 0
+        assert len(pvals_k) == 0
 
 
 def _levy_area_spline():
-    for levy_area in (diffrax.BrownianIncrement, diffrax.SpaceTimeLevyArea):
+    levy_areas = (
+        diffrax.BrownianIncrement,
+        diffrax.SpaceTimeLevyArea,
+        diffrax.SpaceTimeTimeLevyArea,
+    )
+    for levy_area in levy_areas:
         for spline in ("quad", "sqrt", "zero"):
-            if levy_area is diffrax.SpaceTimeLevyArea and spline == "quad":
+            if (
+                issubclass(levy_area, diffrax.AbstractSpaceTimeLevyArea)
+                and spline == "quad"
+            ):
+                # The quad spline is not defined for space-time and
+                # space-time-time Lévy area
                 continue
             yield levy_area, spline
 
 
 @pytest.mark.parametrize("levy_area,spline", _levy_area_spline())
-@pytest.mark.parametrize("use_levy", (False, True))
-def test_spline(levy_area, use_levy, spline):
-    pvals_w1, pvals_w2, pvals_h = conditional_statistics(
+@pytest.mark.parametrize("use_levy", (True, False))
+def test_spline(levy_area, use_levy, spline: _Spline):
+    pvals_w1, pvals_w2, pvals_h, pvals_k, mean_err, cov_err = _conditional_statistics(
         levy_area,
         use_levy=use_levy,
-        tol=2**-4,
-        spacing=2**-5,
+        tol=2**-3,
+        spacing=2**-3,
         spline=spline,
-        min_num_points=50,
+        min_num_points=20,
     )
 
     if spline == "sqrt":
@@ -325,32 +527,153 @@ def test_spline(levy_area, use_levy, spline):
         def pred(pvals):
             return jnp.min(pvals) > 0.1 / pvals.shape[0] and jnp.mean(pvals) > 0.3
 
+        def pred_sttla(_mean_err, _cov_err):
+            return jnp.sum(_mean_err) < 0.005 and jnp.sum(_cov_err) < 0.002
+
     else:
         # make sure that for incorrect splines at least one p-value is
         # below 0.001 (subject to multiple-testing correction) and the
         # average p-value is below 0.03 (i.e. at least 10x smaller than
         # for the correct spline).
         def pred(pvals):
-            return jnp.min(pvals) < 0.001 / pvals.shape[0] and jnp.mean(pvals) < 0.03
+            return jnp.min(pvals) < 0.01 / pvals.shape[0] and jnp.mean(pvals) < 0.03
 
-    if levy_area is diffrax.SpaceTimeLevyArea and use_levy:
+        def pred_sttla(_mean_err, _cov_err):
+            return jnp.sum(_cov_err) > 0.005
+
+    if levy_area is diffrax.SpaceTimeTimeLevyArea and use_levy:
         assert pred(pvals_w2)
         assert pred(pvals_h)
+        assert pred_sttla(mean_err, cov_err)
+
+    elif levy_area is diffrax.SpaceTimeLevyArea and use_levy:
+        assert pred(pvals_w2)
+        assert pred(pvals_h)
+        assert len(pvals_k) == 0
+
     else:
         assert len(pvals_w2) == 0
         assert len(pvals_h) == 0
+        assert len(pvals_k) == 0
         if levy_area is diffrax.BrownianIncrement:
             assert pred(pvals_w1)
         elif spline == "sqrt":  # levy_area == SpaceTimeLevyArea
             assert pred(pvals_w1)
-        else:  # levy_area == SpaceTimeLevyArea and spline == "zero"
+        elif levy_area == diffrax.SpaceTimeLevyArea and spline == "zero":
             # We need a milder upper bound on jnp.mean(pvals_w1) because
-            # the presence of space-time Levy area gives W_r (i.e. the output
+            # the presence of space-time Lévy area gives W_r (i.e. the output
             # of the Brownian path) a variance very close to the correct one,
             # even when the spline is wrong. In pvals_w2 the influence of the
-            # Levy area is subtracted in the mean, so we can use a stricter test.
+            # Lévy area is subtracted in the mean, so we can use a stricter test.
             n = pvals_w1.shape[0]
-            assert jnp.min(pvals_w1) < 0.01 / n and jnp.mean(pvals_w1) < 0.2
+            assert jnp.min(pvals_w1) < 0.03 / n and jnp.mean(pvals_w1) < 0.2
+        else:
+            pass
+
+
+@pytest.mark.parametrize(
+    "tol,spline", [(100.0, "sqrt"), (0.9, "sqrt"), (2**-15, "zero")]
+)
+def test_whk_interpolation(tol, spline):
+    key = jr.key(5678)
+    r_key, bm_key = jr.split(key, 2)
+    s = jnp.array(0.0, dtype=jnp.float64)
+    u = jnp.array(5.7, dtype=jnp.float64)
+    bound = 0.0
+    rs = jr.uniform(
+        r_key, (100,), dtype=jnp.float64, minval=s + bound, maxval=u - bound
+    )
+    path = diffrax.VirtualBrownianTree(
+        t0=s,
+        t1=u,
+        shape=(10000,),
+        tol=tol,
+        key=bm_key,
+        levy_area=diffrax.SpaceTimeTimeLevyArea,
+        _spline=spline,
+    )
+
+    @jax.jit
+    def eval_paths(t):
+        # return jax.vmap(lambda p: p.evaluate(t, use_levy=True))(paths)
+        return path.evaluate(t, use_levy=True)
+
+    bm_s = eval_paths(s)
+    bm_u = eval_paths(u)
+
+    total_mean_err = 0.0
+    total_cov_err = 0.0
+    _pvals_w = []
+    _pvals_h = []
+    _pvals_k = []
+    for r in rs:
+        bm_r = eval_paths(r)
+        w_s = bm_s.W
+        w_r = bm_r.W
+        h_s = bm_s.H
+        h_r = bm_r.H
+        k_s = bm_s.K
+        k_r = bm_r.K
+        sr = r - s
+        sr2 = jnp.square(sr)
+        bh_s = s * h_s
+        bh_r = r * h_r
+        bk_s = s**2 * k_s
+        bk_r = r**2 * k_r
+
+        # Compute the target conditional mean and covariance
+        true_mean_whk, true_cov = _true_cond_stats_whk(bm_s, bm_u, s, r, u)
+
+        # now compute the values of (W_sr, H_sr, K_sr), which are to be tested
+        # against the normal distribution N(mean, cov)
+        w_sr = w_r - w_s
+        r_bb_s = r * w_s - s * w_r
+        bh_sr = bh_r - bh_s - 0.5 * r_bb_s
+        h_sr = bh_sr / sr
+        bk_sr = (
+            bk_r - bk_s - 0.5 * (sr * bh_s - s * bh_sr) - ((r - 2 * s) / 12) * r_bb_s
+        )
+        k_sr = bk_sr / sr2
+
+        y = jnp.stack([w_sr, h_sr, k_sr], axis=0)
+
+        # now we have to confirm that (w_centred, h_centred, k_centred) have
+        # zero mean and covariance matrix cov
+
+        hat_y = y - true_mean_whk
+        tilde_mean = jnp.mean(true_mean_whk, axis=1)
+        y_mean = jnp.mean(y, axis=1)
+        mean_diff = y_mean - tilde_mean
+        emp_cov = jnp.cov(hat_y)
+
+        mean_err = mean_diff
+        cov_err = emp_cov - true_cov
+        total_mean_err += jnp.abs(mean_err) / (len(rs))
+        total_cov_err += jnp.abs(cov_err) / (len(rs))
+
+        hat_w, hat_h, hat_k = hat_y
+        w_var2, h_var, k_var = jnp.diag(true_cov)
+
+        normalised_w2 = hat_w / math.sqrt(w_var2)
+        _, pval_w2 = stats.kstest(normalised_w2, stats.norm.cdf)
+        _pvals_w.append(pval_w2)
+
+        normalised_h = hat_h / math.sqrt(h_var)
+        _, pval_h = stats.kstest(normalised_h, stats.norm.cdf)
+        _pvals_h.append(pval_h)
+
+        normalised_k = hat_k / math.sqrt(k_var)
+        _, pval_k = stats.kstest(normalised_k, stats.norm.cdf)
+        _pvals_k.append(pval_k)
+
+    _pvals_w = jnp.array(_pvals_w)
+    _pvals_h = jnp.array(_pvals_h)
+    _pvals_k = jnp.array(_pvals_k)
+    assert jnp.all(_pvals_w > 0.1 / _pvals_w.shape[0])
+    assert jnp.all(_pvals_h > 0.1 / _pvals_h.shape[0])
+    assert jnp.all(_pvals_k > 0.1 / _pvals_k.shape[0])
+    assert jnp.all(jnp.abs(total_mean_err) < 0.005)
+    assert jnp.all(jnp.abs(total_cov_err) < 0.005)
 
 
 def test_levy_area_reverse_time():
diff --git a/test/test_event.py b/test/test_event.py
index 135ba4de..80f0102c 100644
--- a/test/test_event.py
+++ b/test/test_event.py
@@ -1,8 +1,11 @@
+import functools as ft
 from typing import cast
 
 import diffrax
+import equinox as eqx
 import jax
 import jax.numpy as jnp
+import optimistix as optx
 import pytest
 from jaxtyping import Array
 
@@ -16,13 +19,21 @@ 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)
-    sol = diffrax.diffeqsolve(
-        term, solver, t0, t1, dt0, y0, discrete_terminating_event=event
-    )
+    with pytest.warns(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)
 
 
@@ -35,13 +46,21 @@ 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)
-    sol = diffrax.diffeqsolve(
-        term, solver, t0, t1, dt0, y0, discrete_terminating_event=event
-    )
+    with pytest.warns(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)
 
 
@@ -54,6 +73,7 @@ def test_event_backsolve():
     y0 = 1.0
 
     def event_fn(state, **kwargs):
+        del kwargs
         assert isinstance(state.y, jax.Array)
         return state.tprev > 10
 
@@ -62,21 +82,630 @@ def event_fn(state, **kwargs):
     @jax.jit
     @jax.grad
     def run(y0):
+        with pytest.warns(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.
+    with pytest.raises(NotImplementedError):
+        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()
+    t0 = 0
+    t1 = jnp.inf
+    dt0 = 1
+    y0 = 1.0
+
+    def cond_fn(t, y, args, **kwargs):
+        del y, args, kwargs
+        return t > 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(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)
+    assert jnp.all(cast(Array, sol.ts) >= 10)
+
+
+def test_continuous_event_time():
+    term = diffrax.ODETerm(lambda t, y, args: y)
+    solver = diffrax.Tsit5()
+    t0 = 0
+    t1 = jnp.inf
+    dt0 = 1.0
+    y0 = 1.0
+
+    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)
+    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), 1.0, 1e-4))
+
+
+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(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)
+    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(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)
+    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(t, y, args, **kwargs):
+        del y, args, kwargs
+        assert isinstance(t, jax.Array)
+        return t - 10
+
+    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)
+    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):
+        del t, args
+        _, v = y
+        d_out = v, -8.0
+        return jnp.array(d_out)
+
+    def cond_fn(t, y, args, **kwargs):
+        del t, args, kwargs
+        x, _ = 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(t=None, y=_y, args=None)
+
+        _, 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))
+
+
+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,
-            discrete_terminating_event=event,
-            adjoint=diffrax.BacksolveAdjoint(),
+            stepsize_controller=stepsize_controller,
+            event=event,
         )
-        return jnp.sum(cast(Array, sol.ys))
+        return cast(Array, sol.ys)[-1]
 
-    # And in particular not some other error.
-    with pytest.raises(NotImplementedError):
-        run(y0)
+    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()
 
 
-# diffrax.SteadyStateEvent tested as part of test_adjoint.py::test_implicit
+@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()
+
+
+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(ValueError):
+            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))
diff --git a/test/test_integrate.py b/test/test_integrate.py
index 84d8b3c8..9ed39499 100644
--- a/test/test_integrate.py
+++ b/test/test_integrate.py
@@ -9,6 +9,7 @@
 import jax.numpy as jnp
 import jax.random as jr
 import jax.tree_util as jtu
+import lineax as lx
 import pytest
 import scipy.stats
 from diffrax import ControlTerm, MultiTerm, ODETerm
@@ -644,6 +645,9 @@ class TestSolver(diffrax.AbstractSolver):
             "e": diffrax.MultiTerm[
                 tuple[diffrax.ODETerm, diffrax.AbstractTerm[Any, Float[Array, " 5"]]]
             ],
+            "f": diffrax.MultiTerm[
+                tuple[diffrax.ODETerm, diffrax.AbstractTerm[Any, Float[Array, " 5"]]]
+            ],
         }
         interpolation_cls = diffrax.LocalLinearInterpolation
 
@@ -676,6 +680,13 @@ def func(self, terms, t0, y0, args):
                 lambda t, y, args: -y, lambda t0, t1: jnp.array(t1 - t0).repeat(5)
             ),
         ),
+        "f": diffrax.MultiTerm(
+            ode_term,
+            diffrax.ControlTerm(
+                lambda t, y, args: lx.DiagonalLinearOperator(-y),
+                lambda t0, t1: jnp.array(t1 - t0).repeat(5),
+            ),
+        ),
     }
     compatible_y0 = {
         "a": jnp.array(1.0),
@@ -683,6 +694,7 @@ def func(self, terms, t0, y0, args):
         "c": jnp.arange(3.0),
         "d": jnp.arange(4.0),
         "e": jnp.arange(5.0),
+        "f": jnp.arange(5.0),
     }
     diffrax.diffeqsolve(compatible_term, solver, 0.0, 1.0, 0.1, compatible_y0)
 
@@ -698,6 +710,13 @@ def func(self, terms, t0, y0, args):
                 lambda t0, t1: t1 - t0,  # wrong control shape
             ),
         ),
+        "f": diffrax.MultiTerm(
+            ode_term,
+            diffrax.ControlTerm(
+                lambda t, y, args: lx.DiagonalLinearOperator(-y),
+                lambda t0, t1: jnp.array(t1 - t0).repeat(5),
+            ),
+        ),
     }
     incompatible_term2 = {
         "a": ode_term,
@@ -710,6 +729,13 @@ def func(self, terms, t0, y0, args):
                 lambda t, y, args: -y, lambda t0, t1: jnp.array(t1 - t0).repeat(3)
             ),
         ),
+        "f": diffrax.MultiTerm(
+            ode_term,
+            diffrax.ControlTerm(
+                lambda t, y, args: lx.DiagonalLinearOperator(-y),
+                lambda t0, t1: jnp.array(t1 - t0).repeat(5),
+            ),
+        ),
     }
     incompatible_term3 = {
         "a": ode_term,
@@ -720,6 +746,13 @@ def func(self, terms, t0, y0, args):
         "e": diffrax.WeaklyDiagonalControlTerm(
             lambda t, y, args: -y, lambda t0, t1: jnp.array(t1 - t0).repeat(3)
         ),
+        "f": diffrax.MultiTerm(
+            ode_term,
+            diffrax.ControlTerm(
+                lambda t, y, args: lx.DiagonalLinearOperator(-y),
+                lambda t0, t1: jnp.array(t1 - t0).repeat(5),
+            ),
+        ),
     }
 
     incompatible_y0_1 = {
@@ -728,6 +761,7 @@ def func(self, terms, t0, y0, args):
         "c": jnp.arange(4.0),  # of length 4, not 3
         "d": jnp.arange(4.0),
         "e": jnp.arange(5.0),
+        "f": jnp.arange(5.0),
     }
     incompatible_y0_2 = {
         "a": jnp.array(1.0),
@@ -735,6 +769,7 @@ def func(self, terms, t0, y0, args):
         "c": jnp.arange(3.0),
         # Missing "d" piece
         "e": jnp.arange(5.0),
+        "f": jnp.arange(5.0),
     }
     incompatible_y0_3 = jnp.array(4.0)  # Completely the wrong structure!
     for term in (
diff --git a/test/test_sde.py b/test/test_sde.py
index b2f62245..d44283d8 100644
--- a/test/test_sde.py
+++ b/test/test_sde.py
@@ -5,6 +5,7 @@
 import jax.numpy as jnp
 import jax.random as jr
 import jax.tree_util as jtu
+import lineax as lx
 import pytest
 from diffrax import ControlTerm, MultiTerm, ODETerm, WeaklyDiagonalControlTerm
 
@@ -270,18 +271,63 @@ def _drift(t, y, args):
     assert solution.ys.shape == (1, 3)
 
 
+def _lineax_weakly_diagonal_noise_helper(solver, dtype):
+    w_shape = (3,)
+    args = (0.5, 1.2)
+
+    def _diffusion(t, y, args):
+        a, b = args
+        return lx.DiagonalLinearOperator(jnp.array([b, t, 1 / (t + 1.0)], dtype=dtype))
+
+    def _drift(t, y, args):
+        a, b = args
+        return -a * y
+
+    y0 = jnp.ones(w_shape, dtype)
+
+    bm = diffrax.VirtualBrownianTree(
+        0.0, 1.0, 0.05, w_shape, jr.PRNGKey(0), diffrax.SpaceTimeLevyArea
+    )
+
+    terms = MultiTerm(ODETerm(_drift), ControlTerm(_diffusion, bm))
+    saveat = diffrax.SaveAt(t1=True)
+    solution = diffrax.diffeqsolve(
+        terms, solver, 0.0, 1.0, 0.1, y0, args, saveat=saveat
+    )
+    assert solution.ys is not None
+    assert solution.ys.shape == (1, 3)
+
+
 @pytest.mark.parametrize("solver_ctr", _solvers())
 @pytest.mark.parametrize(
     "dtype",
     (jnp.float64, jnp.complex128),
 )
-def test_weakly_diagonal_noise(solver_ctr, dtype):
-    _weakly_diagonal_noise_helper(solver_ctr(), dtype)
+@pytest.mark.parametrize(
+    "weak_type",
+    ("old", "lineax"),
+)
+def test_weakly_diagonal_noise(solver_ctr, dtype, weak_type):
+    if weak_type == "old":
+        _weakly_diagonal_noise_helper(solver_ctr(), dtype)
+    elif weak_type == "lineax":
+        _lineax_weakly_diagonal_noise_helper(solver_ctr(), dtype)
+    else:
+        raise ValueError("Invalid weak_type")
 
 
 @pytest.mark.parametrize(
     "dtype",
     (jnp.float64, jnp.complex128),
 )
-def test_halfsolver_term_compatible(dtype):
-    _weakly_diagonal_noise_helper(diffrax.HalfSolver(diffrax.SPaRK()), dtype)
+@pytest.mark.parametrize(
+    "weak_type",
+    ("old", "lineax"),
+)
+def test_halfsolver_term_compatible(dtype, weak_type):
+    if weak_type == "old":
+        _weakly_diagonal_noise_helper(diffrax.HalfSolver(diffrax.SPaRK()), dtype)
+    elif weak_type == "lineax":
+        _lineax_weakly_diagonal_noise_helper(diffrax.HalfSolver(diffrax.SPaRK()), dtype)
+    else:
+        raise ValueError("Invalid weak_type")
diff --git a/test/test_term.py b/test/test_term.py
index 420dc122..5260db2c 100644
--- a/test/test_term.py
+++ b/test/test_term.py
@@ -1,8 +1,10 @@
 import diffrax
 import equinox as eqx
+import jax
 import jax.numpy as jnp
 import jax.random as jr
 import jax.tree_util as jtu
+import pytest
 from jaxtyping import Array, PyTree, Shaped
 
 from .helpers import tree_allclose
@@ -45,8 +47,11 @@ def derivative(self, t, left=True):
     y = jnp.array([1.0, 2.0, 3.0])
     vf = term.vf(0, y, args)
     vf_prod = term.vf_prod(0, y, args, dx)
-    assert dx.shape == (2,)
-    assert vf.shape == (3, 2)
+    if isinstance(dx, jax.Array) and isinstance(vf, jax.Array):
+        assert dx.shape == (2,)
+        assert vf.shape == (3, 2)
+    else:
+        raise TypeError("dx/vf is not an array")
     assert vf_prod.shape == (3,)
     assert tree_allclose(vf_prod, term.prod(vf, dx))
 
@@ -85,8 +90,11 @@ def derivative(self, t, left=True):
     y = jnp.array([1.0, 2.0, 3.0])
     vf = term.vf(0, y, args)
     vf_prod = term.vf_prod(0, y, args, dx)
-    assert dx.shape == (3,)
-    assert vf.shape == (3,)
+    if isinstance(dx, jax.Array) and isinstance(vf, jax.Array):
+        assert dx.shape == (3,)
+        assert vf.shape == (3,)
+    else:
+        raise TypeError("dx/vf is not an array")
     assert vf_prod.shape == (3,)
     assert tree_allclose(vf_prod, term.prod(vf, dx))
 
@@ -143,3 +151,10 @@ def __call__(self, t, y, args):
     vf = adjoint_term.vf(t, aug, args)
     vf_prod2 = adjoint_term.prod(vf, dt)
     assert tree_allclose(vf_prod1, vf_prod2)
+
+
+def test_weaklydiagonal_deprecate():
+    with pytest.warns(match="WeaklyDiagonalControlTerm"):
+        _ = diffrax.WeaklyDiagonalControlTerm(
+            lambda t, y, args: 0.0, lambda t0, t1: jnp.array(t1 - t0)
+        )