New sde related fixes

Author: Randl

diffrax/_brownian/tree.py

@@ -91,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
@@ -100,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
@@ -396,27 +397,31 @@ 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:
-            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
+            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
 
@@ -499,8 +504,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/_integrate.py

@@ -155,7 +155,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

diffrax/_solver/implicit_euler.py

@@ -2,7 +2,6 @@
 from typing import ClassVar
 from typing_extensions import TypeAlias
 
-import jax
 import optimistix as optx
 from equinox.internal import ω
 
@@ -82,15 +81,14 @@ def step(
         # write out a `ButcherTableau` and use `AbstractSDIRK`.
         k0 = terms.vf_prod(t0, y0, args, control)
         args = (terms.vf_prod, t1, y0, args, control)
-        with jax.numpy_dtype_promotion("standard"):
-            nonlinear_sol = optx.root_find(
-                _implicit_relation,
-                self.root_finder,
-                k0,
-                args,
-                throw=False,
-                max_steps=self.root_find_max_steps,
-            )
+        nonlinear_sol = optx.root_find(
+            _implicit_relation,
+            self.root_finder,
+            k0,
+            args,
+            throw=False,
+            max_steps=self.root_find_max_steps,
+        )
         k1 = nonlinear_sol.value
         y1 = (y0**ω + k1**ω).ω
         # Use the trapezoidal rule for adaptive step sizing.

diffrax/_solver/srk.py

@@ -12,6 +12,7 @@
 import numpy as np
 from equinox.internal import ω
 from jaxtyping import Array, Float, PyTree
+from lineax.internal import complex_to_real_dtype
 
 from .._custom_types import (
     AbstractBrownianIncrement,
@@ -340,7 +341,9 @@ def step(
 
         # First the drift related stuff
         a = self._embed_a_lower(self.tableau.a, dtype)
-        c = jnp.asarray(np.insert(self.tableau.c, 0, 0.0), dtype=dtype)
+        c = jnp.asarray(
+            np.insert(self.tableau.c, 0, 0.0), dtype=complex_to_real_dtype(dtype)
+        )
         b_sol = jnp.asarray(self.tableau.b_sol, dtype=dtype)
 
         def make_zeros():

test/helpers.py

@@ -18,7 +18,6 @@
 )
 from jax import Array
 from jaxtyping import PRNGKeyArray, PyTree, Shaped
-from lineax.internal import complex_to_real_dtype
 
 
 all_ode_solvers = (
@@ -252,7 +251,6 @@ def sde_solver_strong_order(
         bm_tol,
         saveat,
     )
-    dts = 2.0 ** jnp.arange(-3, -3 - num_levels, -1, dtype=dtype)
 
     errs_list, steps_list = [], []
     for level in range(level_coarse, level_fine + 1):
@@ -277,7 +275,8 @@ def sde_solver_strong_order(
         steps_list.append(jnp.average(steps))
     errs_arr = jnp.array(errs_list)
     steps_arr = jnp.array(steps_list)
-    order, _ = jnp.polyfit(jnp.log(1 / steps_arr), jnp.log(errs_arr), 1)
+    with jax.numpy_dtype_promotion("standard"):
+        order, _ = jnp.polyfit(jnp.log(1 / steps_arr), jnp.log(errs_arr), 1)
     return steps_arr, errs_arr, order
 
 
@@ -360,12 +359,14 @@ def _squareplus(x):
 
 def drift(t, y, args):
     mlp, _, _ = args
-    return 0.25 * mlp(y)
+    with jax.numpy_dtype_promotion("standard"):
+        return 0.25 * mlp(y)
 
 
 def diffusion(t, y, args):
     _, mlp, noise_dim = args
-    return 1.0 * mlp(y).reshape(3, noise_dim)
+    with jax.numpy_dtype_promotion("standard"):
+        return 1.0 * mlp(y).reshape(3, noise_dim)
 
 
 def get_mlp_sde(t0, t1, dtype, key, noise_dim):
@@ -447,8 +448,9 @@ def ft(t):
     drift_mlp = init_linear_weight(drift_mlp, lap_init, driftkey)
 
     def _drift(t, y, _):
-        mlp_out = drift_mlp(jnp.concatenate([y, ft(t)]))
-        return (0.01 * mlp_out - 0.5 * y**3) / (jnp.sum(y**2) + 1)
+        with jax.numpy_dtype_promotion("standard"):
+            mlp_out = drift_mlp(jnp.concatenate([y, ft(t)]))
+            return (0.01 * mlp_out - 0.5 * y**3) / (jnp.sum(y**2) + 1)
 
     diffusion_mx = jr.normal(diffusionkey, (4, y_dim, noise_dim), dtype=dtype)
 

test/test_sde.py

@@ -43,8 +43,12 @@ def _solvers_and_orders():
 # converges to its own limit (i.e. using itself as reference), and then in a
 # different test check whether that limit is the same as the Euler/Heun limit.
 @pytest.mark.parametrize("solver_ctr,noise,theoretical_order", _solvers_and_orders())
+@pytest.mark.parametrize(
+    "dtype",
+    (jnp.float64,),
+)
 def test_sde_strong_order_new(
-    solver_ctr, noise: Literal["any", "com", "add"], theoretical_order
+    solver_ctr, noise: Literal["any", "com", "add"], theoretical_order, dtype
 ):
     bmkey = jr.PRNGKey(5678)
     sde_key = jr.PRNGKey(11)
@@ -54,15 +58,15 @@ def test_sde_strong_order_new(
     t1 = 5.3
 
     if noise == "add":
-        sde = get_time_sde(t0, t1, jnp.float64, sde_key, noise_dim=7)
+        sde = get_time_sde(t0, t1, dtype, sde_key, noise_dim=7)
     else:
         if noise == "com":
             noise_dim = 1
         elif noise == "any":
             noise_dim = 5
         else:
             assert False
-        sde = get_mlp_sde(t0, t1, jnp.float64, sde_key, noise_dim=noise_dim)
+        sde = get_mlp_sde(t0, t1, dtype, sde_key, noise_dim=noise_dim)
 
     ref_solver = solver_ctr()
     level_coarse, level_fine = 1, 7

test/test_solver.py

@@ -274,7 +274,7 @@ def order(self, terms):
 
 
 # Essentially used as a check that our general IMEX implementation is correct.
-@pytest.mark.parametrize("dtype", (jnp.float64, jnp.complex128))
+@pytest.mark.parametrize("dtype", (jnp.float64,))
 def test_sil3(dtype):
     class ReferenceSil3(diffrax.AbstractImplicitSolver):
         term_structure = diffrax.MultiTerm[
@@ -314,8 +314,7 @@ def _second_stage(ya, _):
                 return ya - (y0 + (1 / 3) * f0 + (1 / 6) * g0 + (1 / 6) * g1)
 
             ta = t0 + (1 / 3) * dt
-            with jax.numpy_dtype_promotion("standard"):
-                ya = optx.root_find(_second_stage, self.root_finder, y0).value
+            ya = optx.root_find(_second_stage, self.root_finder, y0).value
             fs.append(ex_vf_prod(ta, ya))
             gs.append(im_vf_prod(ta, ya))
 
@@ -329,8 +328,7 @@ def _third_stage(yb, _):
 
             tb = t0 + (2 / 3) * dt
 
-            with jax.numpy_dtype_promotion("standard"):
-                yb = optx.root_find(_third_stage, self.root_finder, ya).value
+            yb = optx.root_find(_third_stage, self.root_finder, ya).value
             fs.append(ex_vf_prod(tb, yb))
             gs.append(im_vf_prod(tb, yb))
 
@@ -349,8 +347,7 @@ def _fourth_stage(yc, _):
                 )
 
             tc = t1
-            with jax.numpy_dtype_promotion("standard"):
-                yc = optx.root_find(_fourth_stage, self.root_finder, yb).value
+            yc = optx.root_find(_fourth_stage, self.root_finder, yb).value
             fs.append(ex_vf_prod(tc, yc))
             gs.append(im_vf_prod(tc, yc))