Merge remote-tracking branch 'upstream/main'

Author: lockwo

benchmarks/brownian_tree_times.py

@@ -10,6 +10,7 @@
 from typing import cast, Optional, Union
 from typing_extensions import TypeAlias
 
+import diffrax
 import equinox as eqx
 import equinox.internal as eqxi
 import jax
@@ -50,9 +51,9 @@ def __init__(
         tol: RealScalarLike,
         shape: tuple[int, ...],
         key: PRNGKeyArray,
-        levy_area: str,
+        levy_area: type[diffrax.AbstractBrownianIncrement] = diffrax.BrownianIncrement,
     ):
-        assert levy_area == ""
+        assert levy_area == diffrax.BrownianIncrement
         self.t0 = t0
         self.t1 = t1
         self.tol = tol
@@ -187,13 +188,13 @@ def run(_ts):
     )
 
 
-for levy_area in ("", "space-time"):
+for levy_area in (diffrax.BrownianIncrement, diffrax.SpaceTimeLevyArea):
     print(f"-   {levy_area=}")
     for tol in (2**-3, 2**-12):
         print(f"--  {tol=}")
-        for num_ts in (1, 100):
+        for num_ts in (1, 10000):
             print(f"--- {num_ts=}")
-            if levy_area == "":
+            if levy_area == diffrax.BrownianIncrement:
                 print(f"Old: {time_tree(OldVBT, num_ts, tol, levy_area):.5f}")
             print(f"new: {time_tree(VirtualBrownianTree, num_ts, tol, levy_area):.5f}")
     print("")

diffrax/_adjoint.py

@@ -429,6 +429,14 @@ def _solve(inputs):
     )
 
 
+# Unwrap jaxtyping decorator during tests, so that these are global functions.
+# This is needed to ensure `optx.implicit_jvp` is happy.
+if _vf.__globals__["__name__"].startswith("jaxtyping"):
+    _vf = _vf.__wrapped__  # pyright: ignore[reportFunctionMemberAccess]
+if _solve.__globals__["__name__"].startswith("jaxtyping"):
+    _solve = _solve.__wrapped__  # pyright: ignore[reportFunctionMemberAccess]
+
+
 def _frozenset(x: Union[object, Iterable[object]]) -> frozenset[object]:
     try:
         iter_x = iter(x)  # pyright: ignore

diffrax/_brownian/path.py

@@ -9,6 +9,7 @@
 import jax.tree_util as jtu
 import lineax.internal as lxi
 from jaxtyping import Array, PRNGKeyArray, PyTree
+from lineax.internal import complex_to_real_dtype
 
 from .._custom_types import (
     AbstractBrownianIncrement,
@@ -44,7 +45,19 @@ class UnsafeBrownianPath(AbstractBrownianPath):
     motion. Hence the restrictions above. (They describe the general case for which the
     correlation structure isn't needed.)
 
-    Depending on the `levy_area` argument, this can also be used to generate Levy area.
+    !!! info "Levy Area"
+
+        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.
+
+        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`.
     """
 
     shape: PyTree[jax.ShapeDtypeStruct] = eqx.field(static=True)
@@ -130,7 +143,7 @@ def _evaluate_leaf(
     ):
         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=shape.dtype)
+        dt = jnp.asarray(t1 - t0, dtype=complex_to_real_dtype(shape.dtype))
 
         if levy_area is SpaceTimeLevyArea:
             key, key_hh = jr.split(key, 2)

diffrax/_brownian/tree.py

@@ -10,7 +10,8 @@
 import jax.random as jr
 import jax.tree_util as jtu
 import lineax.internal as lxi
-from jaxtyping import Array, Float, PRNGKeyArray, PyTree
+from jaxtyping import Array, Inexact, PRNGKeyArray, PyTree
+from lineax.internal import complex_to_real_dtype
 
 from .._custom_types import (
     AbstractBrownianIncrement,
@@ -54,9 +55,9 @@
 # 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.
 
-FloatDouble: TypeAlias = tuple[Float[Array, " *shape"], Float[Array, " *shape"]]
+FloatDouble: TypeAlias = tuple[Inexact[Array, " *shape"], Inexact[Array, " *shape"]]
 FloatTriple: TypeAlias = tuple[
-    Float[Array, " *shape"], Float[Array, " *shape"], Float[Array, " *shape"]
+    Inexact[Array, " *shape"], Inexact[Array, " *shape"], Inexact[Array, " *shape"]
 ]
 _Spline: TypeAlias = Literal["sqrt", "quad", "zero"]
 _BrownianReturn = TypeVar("_BrownianReturn", bound=AbstractBrownianIncrement)
@@ -90,7 +91,7 @@ def _levy_diff(_, x0: tuple, x1: tuple) -> AbstractBrownianIncrement:
         assert len(x1) == 2
         dt0, w0 = x0
         dt1, w1 = x1
-        su = jnp.asarray(dt1 - dt0, dtype=w0.dtype)
+        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
@@ -99,12 +100,13 @@ def _levy_diff(_, x0: tuple, x1: tuple) -> AbstractBrownianIncrement:
         dt1, w1, hh1, bhh1 = x1
 
         w_su = w1 - w0
-        su = jnp.asarray(dt1 - dt0, dtype=w0.dtype)
+        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
-        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
+        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)
     else:
         assert False
@@ -135,10 +137,19 @@ def _split_interval(
 class VirtualBrownianTree(AbstractBrownianPath):
     """Brownian simulation that discretises the interval `[t0, t1]` to tolerance `tol`.
 
-    Can be initialised with `levy_area` set to `""`, or `"space-time"`.
-    If `levy_area="space_time"`, then it also computes space-time Lévy area `H`.
-    This will impact the Brownian path, so even with the same key, the trajectory will
-    be different depending on the value of `levy_area`.
+    !!! info "Levy Area"
+
+        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.
+
+        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`.
 
     ??? cite "Reference"
 
@@ -283,9 +294,10 @@ def _evaluate_leaf(
         tuple[RealScalarLike, Array], tuple[RealScalarLike, Array, Array, Array]
     ]:
         shape, dtype = struct.shape, struct.dtype
+        tdtype = complex_to_real_dtype(dtype)
 
-        t0 = jnp.zeros((), dtype)
-        r = jnp.asarray(r, 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)
@@ -394,14 +406,33 @@ def _body_fun(_state: _State):
             a = d_prime * sr3 * sr_ru_half
             b = d_prime * ru3 * sr_ru_half
 
-            w_sr = sr / su * w_su + 6 * sr * ru / su3 * bhh_su + 2 * (a + b) / su * x1
-            w_r = w_s + w_sr
-            c = jnp.sqrt(3 * sr3 * ru3) / (6 * d)
-            bhh_sr = sr3 / su3 * bhh_su - a * x1 + c * x2
-            bhh_r = bhh_s + bhh_sr + 0.5 * (r * w_s - s * w_r)
+            with jax.numpy_dtype_promotion("standard"):
+                w_sr = (
+                    sr / su * w_su + 6 * sr * ru / su3 * bhh_su + 2 * (a + b) / su * x1
+                )
+                w_r = w_s + w_sr
+                c = jnp.sqrt(3 * sr3 * ru3) / (6 * d)
+                bhh_sr = sr3 / su3 * bhh_su - a * x1 + c * x2
+                bhh_r = bhh_s + bhh_sr + 0.5 * (r * w_s - s * w_r)
 
-            inverse_r = 1 / jnp.where(jnp.abs(r) < jnp.finfo(r).eps, jnp.inf, r)
-            hh_r = inverse_r * bhh_r
+                inverse_r = 1 / jnp.where(jnp.abs(r) < jnp.finfo(r).eps, jnp.inf, r)
+                hh_r = inverse_r * bhh_r
+
+        elif self.levy_area is BrownianIncrement:
+            with jax.numpy_dtype_promotion("standard"):
+                w_mean = w_s + sr / su * w_su
+                if self._spline == "sqrt":
+                    z = jr.normal(final_state.key, shape, dtype)
+                    bb = jnp.sqrt(sr * ru / su) * z
+                elif self._spline == "quad":
+                    z = jr.normal(final_state.key, shape, dtype)
+                    bb = (sr * ru / su) * z
+                elif self._spline == "zero":
+                    bb = jnp.zeros(shape, dtype)
+                else:
+                    assert False
+            w_r = w_mean + bb
+            return r, w_r
 
         elif self.levy_area is BrownianIncrement:
             w_mean = w_s + sr / su * w_su
@@ -497,8 +528,8 @@ def _brownian_arch(
 
             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)
+            with jax.numpy_dtype_promotion("standard"):
+                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

diffrax/_global_interpolation.py

@@ -143,7 +143,10 @@ def _index(_ys):
         next_t = self.ts[index + 1]
         diff_t = next_t - prev_t
 
-        return (prev_ys**ω + (next_ys**ω - prev_ys**ω) * (fractional_part / diff_t)).ω
+        with jax.numpy_dtype_promotion("standard"):
+            return (
+                prev_ys**ω + (next_ys**ω - prev_ys**ω) * (fractional_part / diff_t)
+            ).ω
 
     @eqx.filter_jit
     def derivative(self, t: RealScalarLike, left: bool = True) -> PyTree[Array]:
@@ -165,10 +168,11 @@ def derivative(self, t: RealScalarLike, left: bool = True) -> PyTree[Array]:
 
         index, _ = self._interpret_t(t, left)
 
-        return (
-            (ω(self.ys)[index + 1] - ω(self.ys)[index])
-            / (self.ts[index + 1] - self.ts[index])
-        ).ω
+        with jax.numpy_dtype_promotion("standard"):
+            return (
+                (ω(self.ys)[index + 1] - ω(self.ys)[index])
+                / (self.ts[index + 1] - self.ts[index])
+            ).ω
 
 
 LinearInterpolation.__init__.__doc__ = """**Arguments:**
@@ -254,10 +258,11 @@ def evaluate(
 
         d, c, b, a = self.coeffs
 
-        return (
-            ω(a)[index]
-            + frac * (ω(b)[index] + frac * (ω(c)[index] + frac * ω(d)[index]))
-        ).ω
+        with jax.numpy_dtype_promotion("standard"):
+            return (
+                ω(a)[index]
+                + frac * (ω(b)[index] + frac * (ω(c)[index] + frac * ω(d)[index]))
+            ).ω
 
     @eqx.filter_jit
     def derivative(
@@ -283,7 +288,8 @@ def derivative(
 
         d, c, b, _ = self.coeffs
 
-        return (ω(b)[index] + frac * (2 * ω(c)[index] + frac * 3 * ω(d)[index])).ω
+        with jax.numpy_dtype_promotion("standard"):
+            return (ω(b)[index] + frac * (2 * ω(c)[index] + frac * 3 * ω(d)[index])).ω
 
 
 CubicInterpolation.__init__.__doc__ = """**Arguments:**
@@ -622,8 +628,9 @@ def _hermite_forward(
 ]:
     prev_ti, prev_yi, prev_deriv_i = carry
     ti, yi, next_ti, next_yi = value
-    first_deriv_i = (next_yi - yi) / (next_ti - ti)
-    later_deriv_i = (yi - prev_yi) / (ti - prev_ti)
+    with jax.numpy_dtype_promotion("standard"):
+        first_deriv_i = (next_yi - yi) / (next_ti - ti)
+        later_deriv_i = (yi - prev_yi) / (ti - prev_ti)
     deriv_i = jnp.where(jnp.isnan(prev_yi), first_deriv_i, later_deriv_i)
     cond = jnp.isnan(yi)
     carry_ti = jnp.where(cond, prev_ti, ti)
@@ -635,13 +642,15 @@ def _hermite_forward(
 
 def _hermite_coeffs(t0, y0, deriv0, t1, y1):
     t_diff = t1 - t0
-    deriv1 = (y1 - y0) / t_diff
-    d_deriv = deriv1 - deriv0
 
-    a = y0
-    b = deriv0
-    c = 2 * d_deriv / t_diff
-    d = -d_deriv / t_diff**2
+    with jax.numpy_dtype_promotion("standard"):
+        deriv1 = (y1 - y0) / t_diff
+        d_deriv = deriv1 - deriv0
+
+        a = y0
+        b = deriv0
+        c = 2 * d_deriv / t_diff
+        d = -d_deriv / (t_diff**2)
 
     return d, c, b, a
 
@@ -684,7 +693,8 @@ def _backward_hermite_coefficients(
     else:
         y0 = jnp.broadcast_to(replace_nans_at_start, ys[0].shape)
     if deriv0 is None:
-        deriv0 = (next_ys[0] - y0) / (next_ts[0] - t0)
+        with jax.numpy_dtype_promotion("standard"):
+            deriv0 = (next_ys[0] - y0) / (next_ts[0] - t0)
     else:
         deriv0 = jnp.broadcast_to(deriv0, ys[0].shape)
     ts = ts[:-1]

diffrax/_integrate.py

@@ -156,7 +156,8 @@ def _check(term_cls, term, term_contr_kwargs, yi):
             # If we've got to this point then the term is compatible
 
     try:
-        jtu.tree_map(_check, term_structure, terms, contr_kwargs, y)
+        with jax.numpy_dtype_promotion("standard"):
+            jtu.tree_map(_check, term_structure, terms, contr_kwargs, y)
     except ValueError:
         # ValueError may also arise from mismatched tree structures
         return False
@@ -709,11 +710,6 @@ def diffeqsolve(
         eqx.is_array(xi) and jnp.iscomplexobj(xi)
         for xi in jtu.tree_leaves((terms, y0, args))
     ):
-        if isinstance(solver, AbstractImplicitSolver):
-            raise ValueError(
-                "Implicit solvers in conjunction with complex dtypes is currently not "
-                "supported."
-            )
         warnings.warn(
             "Complex dtype support is work in progress, please read "
             "https://github.com/patrick-kidger/diffrax/pull/197 and proceed carefully.",
@@ -808,9 +804,6 @@ def _promote(yi):
                 "`UnsafeBrownianPath` cannot be used with adaptive step sizes."
             )
 
-    if isinstance(solver, KLSolver):
-        y0 = (y0, 0.0)
-        y0 = jtu.tree_map(_promote, y0)
     # Normalises time: if t0 > t1 then flip things around.
     direction = jnp.where(t0 < t1, 1, -1)
     t0 = t0 * direction