New sde related fixes

Author: Randl

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,
@@ -130,7 +131,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/_solver/srk.py

@@ -406,7 +406,7 @@ def aux_add_levy(w_leaf, *levy_leaves):
             def _comp_g(_t):
                 return diffusion.vf(_t, y0, args)
 
-            g0_g1 = _comp_g(jnp.array([t0, t1], dtype=dtype))
+            g0_g1 = _comp_g(jnp.array([t0, t1], dtype=complex_to_real_dtype(dtype)))
             g0 = jtu.tree_map(lambda g_leaf: g_leaf[0], g0_g1)
             # g_delta = 0.5 * g1 - g0
             g_delta = jtu.tree_map(lambda g_leaf: 0.5 * (g_leaf[1] - g_leaf[0]), g0_g1)
@@ -534,13 +534,15 @@ def compute_and_insert_kf_j(_h_kfs_in):
                 return (_h_kfs, None, None), None
 
             def compute_and_insert_kg_j(_w_kgs_in, _levylist_kgs_in):
-                _w_kg_j = diffusion.vf_prod(t0 + c_j * h, z_j, args, w)
+                with jax.numpy_dtype_promotion("standard"):
+                    _w_kg_j = diffusion.vf_prod(t0 + c_j * h, z_j, args, w)
                 new_w_kgs = insert_jth_stage(_w_kgs_in, _w_kg_j, j)
 
-                _levylist_kg_j = [
-                    diffusion.vf_prod(t0 + c_j * h, z_j, args, levy)
-                    for levy in levy_areas
-                ]
+                with jax.numpy_dtype_promotion("standard"):
+                    _levylist_kg_j = [
+                        diffusion.vf_prod(t0 + c_j * h, z_j, args, levy)
+                        for levy in levy_areas
+                    ]
                 new_levylist_kgs = insert_jth_stage(_levylist_kgs_in, _levylist_kg_j, j)
                 return new_w_kgs, new_levylist_kgs
 

diffrax/_term.py

@@ -361,7 +361,8 @@ class WeaklyDiagonalControlTerm(_AbstractControlTerm[_VF, _Control]):
     """
 
     def prod(self, vf: _VF, control: _Control) -> Y:
-        return jtu.tree_map(operator.mul, vf, control)
+        with jax.numpy_dtype_promotion("standard"):
+            return jtu.tree_map(operator.mul, vf, control)
 
 
 class _ControlToODE(eqx.Module):
@@ -461,7 +462,8 @@ def contr(self, t0: RealScalarLike, t1: RealScalarLike, **kwargs) -> _Control:
         return (self.direction * self.term.contr(_t0, _t1, **kwargs) ** ω).ω
 
     def prod(self, vf: _VF, control: _Control) -> Y:
-        return self.term.prod(vf, control)
+        with jax.numpy_dtype_promotion("standard"):
+            return self.term.prod(vf, control)
 
     def vf_prod(self, t: RealScalarLike, y: Y, args: Args, control: _Control) -> Y:
         t = t * self.direction

test/helpers.py

@@ -99,7 +99,8 @@ def path_l2_dist(
     # all but the first two axes (which represent the number of samples
     # and the length of saveat). Also sum all the PyTree leaves.
     def sum_square_diff(y1, y2):
-        square_diff = jnp.square(y1 - y2)
+        with jax.numpy_dtype_promotion("standard"):
+            square_diff = jnp.square(y1 - y2)
         # sum all but the first two axes
         axes = range(2, square_diff.ndim)
         out = jnp.sum(square_diff, axis=axes)

test/test_brownian.py

@@ -46,6 +46,11 @@ def test_shape_and_dtype(ctr, levy_area, use_levy, getkey):
         (2,),
         (3, 4),
         (1, 2, 3, 4),
+        (1, 2, 3, 4),
+        {
+            "a": (1,),
+            "b": (2, 3),
+        },
         {
             "a": (1,),
             "b": (2, 3),
@@ -66,7 +71,9 @@ def test_shape_and_dtype(ctr, levy_area, use_levy, getkey):
         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)),
     )
 

test/test_sde.py

@@ -116,8 +116,12 @@ def get_dt_and_controller(level):
 # using a single reference solution. We use Euler if the solver is Ito
 # and Heun if the solver is Stratonovich.
 @pytest.mark.parametrize("solver_ctr,noise,theoretical_order", _solvers_and_orders())
+@pytest.mark.parametrize(
+    "dtype",
+    (jnp.float64,),
+)
 def test_sde_strong_limit(
-    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)
@@ -127,7 +131,7 @@ def test_sde_strong_limit(
     t1 = 5.3
 
     if noise == "add":
-        sde = get_time_sde(t0, t1, jnp.float64, sde_key, noise_dim=3)
+        sde = get_time_sde(t0, t1, dtype, sde_key, noise_dim=3)
         level_fine = 12
         if theoretical_order <= 1.0:
             level_coarse = 11
@@ -141,7 +145,7 @@ def test_sde_strong_limit(
             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)
 
     # Reference solver is always an ODE-viable solver, so its implementation has been
     # verified by the ODE tests like test_ode_order.
@@ -210,9 +214,12 @@ def get_matrix(y_leaf):
 
 @pytest.mark.parametrize("shape", [(), (5, 2)])
 @pytest.mark.parametrize("solver_ctr", _solvers())
-def test_sde_solver_shape(shape, solver_ctr):
+@pytest.mark.parametrize(
+    "dtype",
+    (jnp.float64, jnp.complex128),
+)
+def test_sde_solver_shape(shape, solver_ctr, dtype):
     pytree = ({"a": 0, "b": [0, 0]}, 0, 0)
-    dtype = jnp.float64
     key = jr.PRNGKey(0)
     y0 = jtu.tree_map(lambda _: jr.normal(key, shape, dtype=dtype), pytree)
     t0, t1, dt0 = 0.0, 1.0, 0.3
@@ -236,8 +243,7 @@ def test_sde_solver_shape(shape, solver_ctr):
         assert leaf[0].shape == shape
 
 
-def _weakly_diagonal_noise_helper(solver):
-    dtype = jnp.float64
+def _weakly_diagonal_noise_helper(solver, dtype):
     w_shape = (3,)
     args = (0.5, 1.2)
 
@@ -265,9 +271,17 @@ def _drift(t, y, args):
 
 
 @pytest.mark.parametrize("solver_ctr", _solvers())
-def test_weakly_diagonal_noise(solver_ctr):
-    _weakly_diagonal_noise_helper(solver_ctr())
+@pytest.mark.parametrize(
+    "dtype",
+    (jnp.float64, jnp.complex128),
+)
+def test_weakly_diagonal_noise(solver_ctr, dtype):
+    _weakly_diagonal_noise_helper(solver_ctr(), dtype)
 
 
-def test_halfsolver_term_compatible():
-    _weakly_diagonal_noise_helper(diffrax.HalfSolver(diffrax.SPaRK()))
+@pytest.mark.parametrize(
+    "dtype",
+    (jnp.float64, jnp.complex128),
+)
+def test_halfsolver_term_compatible(dtype):
+    _weakly_diagonal_noise_helper(diffrax.HalfSolver(diffrax.SPaRK()), dtype)