Merge branch 'patrick-kidger:main' into main

Author: thibmonsel

diffrax/integrate.py

@@ -19,6 +19,7 @@
 from .saveat import SaveAt, SubSaveAt
 from .solution import is_okay, is_successful, RESULTS, Solution
 from .solver import (
+    AbstractImplicitSolver,
     AbstractItoSolver,
     AbstractSolver,
     AbstractStratonovichSolver,
@@ -605,6 +606,18 @@ def diffeqsolve(
             pred = (t1 - t0) * dt0 < 0
         dt0 = eqxi.error_if(jnp.array(dt0), pred, msg)
 
+    # Error checking and warning for complex dtypes
+    if any(jtu.tree_leaves(jtu.tree_map(jnp.iscomplexobj, y0))):
+        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."
+        )
+
     # Backward compatibility
     if isinstance(
         solver, (EulerHeun, ItoMilstein, StratonovichMilstein)
@@ -664,8 +677,10 @@ def _get_subsaveat_ts(saveat):
     )
 
     # Time will affect state, so need to promote the state dtype as well if necessary.
+    # fixing issue with float64 and weak dtypes, see discussion at:
+    # https://github.com/patrick-kidger/diffrax/pull/197#discussion_r1130173527
     def _promote(yi):
-        _dtype = jnp.result_type(yi, *timelikes)  # noqa: F821
+        _dtype = jnp.result_type(yi, dtype)  # noqa: F821
         return jnp.asarray(yi, dtype=_dtype)
 
     y0 = jtu.tree_map(_promote, y0)
@@ -759,7 +774,9 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
         save_index = 0
         ts = jnp.full(out_size, jnp.inf)
         struct = eqx.filter_eval_shape(subsaveat.fn, t0, y0, args)
-        ys = jtu.tree_map(lambda y: jnp.full((out_size,) + y.shape, jnp.inf), struct)
+        ys = jtu.tree_map(
+            lambda y: jnp.full((out_size,) + y.shape, jnp.inf, dtype=y.dtype), struct
+        )
         return SaveState(
             ts=ts, ys=ys, save_index=save_index, saveat_ts_index=saveat_ts_index
         )
@@ -779,7 +796,9 @@ def _allocate_output(subsaveat: SubSaveAt) -> SaveState:
             solver.step, terms, tprev, tnext, y0, args, solver_state, made_jump
         )
         dense_ts = jnp.full(max_steps + 1, jnp.inf)
-        _make_full = lambda x: jnp.full((max_steps,) + jnp.shape(x), jnp.inf)
+        _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_save_index = 0
     else:

diffrax/misc.py

@@ -107,8 +107,8 @@ def _rms_norm_jvp(x, tx):
     pred = (out == 0) | jnp.isinf(out)
     numerator = jnp.where(pred, 0, x)
     denominator = jnp.where(pred, 1, out * x.size)
-    t_out = jnp.dot(numerator / denominator, tx)
-    return out, t_out
+    t_out = jnp.dot(numerator / denominator, jnp.conj(tx))
+    return out, jnp.real(t_out)
 
 
 def adjoint_rms_seminorm(x: Tuple[PyTree, PyTree, PyTree, PyTree]) -> Scalar:

diffrax/nonlinear_solver/base.py

@@ -106,4 +106,6 @@ def jac(fn: Callable, x: PyTree, args: PyTree) -> LU_Jacobian:
         if not jnp.issubdtype(flat, jnp.inexact):
             # Handle integer arguments
             flat = flat.astype(jnp.float32)
-        return jsp.linalg.lu_factor(jax.jacfwd(curried)(flat))
+        return jsp.linalg.lu_factor(
+            jax.jacfwd(curried, holomorphic=jnp.iscomplexobj(flat))(flat)
+        )

diffrax/step_size_controller/adaptive.py

@@ -424,7 +424,7 @@ def adapt_step_size(
         # ε_n     = atol + norm(y) * rtol with y on the nth step
         # r_n     = norm(y_error) with y_error on the nth step
         # δ_{n,m} = norm(y_error / (atol + norm(y) * rtol))^(-1) with y_error on the nth
-        #                                                        step and y on the mth step
+        # step and y on the mth step
         # β_1     = pcoeff + icoeff + dcoeff
         # β_2     = -(pcoeff + 2 * dcoeff)
         # β_3     = dcoeff

examples/hessian.ipynb

@@ -78,9 +78,9 @@
    "id": "a3ec6532-5b0a-4e4c-af33-bef58c0a7319",
    "metadata": {},
    "source": [
-    "Note the use of the `scan_kind` argument to `Tsit5`. By default, Diffrax internally uses constructs that are optimised specifically for first-order reverse mode autodifferentiation. This argument is needed to switch to a different implementation that is compatible with higher-order autodiff. (In this case: for the loop-over-stages in the Runge--Kutta solver.)\n",
+    "Note the use of the `scan_kind` argument to `Tsit5`. By default, Diffrax internally uses constructs that are optimised specifically for first-order reverse-mode autodifferentiation. This argument is needed to switch to a different implementation that is compatible with higher-order autodiff. (In this case: for the loop-over-stages in the Runge--Kutta solver.)\n",
     "\n",
-    "In similar fashion, if using `saveat=SaveAt(steps=True)` then you will need to pass `adjoint=DirectAdjoint()`. (In this case: for the loop-over-saving output.)"
+    "In similar fashion, if using `saveat=SaveAt(ts=...)` (or a handful of other esoteric cases) then you will need to pass `adjoint=DirectAdjoint()`. (In this case: for the loop-over-saving output.)"
    ]
   }
  ],

test/helpers.py

@@ -44,14 +44,14 @@ def implicit_tol(solver):
     return solver
 
 
-def random_pytree(key, treedef):
+def random_pytree(key, treedef, dtype=None):
     keys = jrandom.split(key, treedef.num_leaves)
     leaves = []
     for key in keys:
         dimkey, sizekey, valuekey = jrandom.split(key, 3)
         num_dims = jrandom.randint(dimkey, (), 0, 5)
         dim_sizes = jrandom.randint(sizekey, (num_dims,), 0, 5)
-        value = jrandom.normal(valuekey, dim_sizes)
+        value = jrandom.normal(valuekey, dim_sizes, dtype=dtype)
         leaves.append(value)
     return jtu.tree_unflatten(treedef, leaves)
 

test/test_integrate.py

@@ -41,7 +41,7 @@ def _all_pairs(*args):
 
 
 @pytest.mark.parametrize(
-    "solver,t_dtype,treedef,stepsize_controller",
+    "solver,t_dtype,y_dtype,treedef,stepsize_controller",
     _all_pairs(
         dict(
             default=diffrax.Euler(),
@@ -58,21 +58,32 @@ def _all_pairs(*args):
             ),
         ),
         dict(default=jnp.float32, opts=(int, float, jnp.int32)),
+        dict(default=jnp.float32, opts=(jnp.complex64,)),
         dict(default=treedefs[0], opts=treedefs[1:]),
         dict(
             default=diffrax.ConstantStepSize(),
             opts=(diffrax.PIDController(rtol=1e-3, atol=1e-6),),
         ),
     ),
 )
-def test_basic(solver, t_dtype, treedef, stepsize_controller, getkey):
+def test_basic(solver, t_dtype, y_dtype, treedef, stepsize_controller, getkey):
     if not isinstance(solver, diffrax.AbstractAdaptiveSolver) and isinstance(
         stepsize_controller, diffrax.PIDController
     ):
         return
 
-    def f(t, y, args):
-        return jtu.tree_map(operator.neg, y)
+    if jnp.iscomplexobj(y_dtype):
+
+        def f(t, y, args):
+            return jtu.tree_map(lambda _y: operator.mul(-1j, _y), y)
+
+        if isinstance(solver, diffrax.AbstractImplicitSolver):
+            return
+
+    else:
+
+        def f(t, y, args):
+            return jtu.tree_map(operator.neg, y)
 
     if t_dtype is int:
         t0 = 0
@@ -92,7 +103,7 @@ def f(t, y, args):
         dt0 = jnp.array(0.01)
     else:
         raise ValueError
-    y0 = random_pytree(getkey(), treedef)
+    y0 = random_pytree(getkey(), treedef, dtype=y_dtype)
     try:
         sol = diffrax.diffeqsolve(
             diffrax.ODETerm(f),
@@ -113,7 +124,10 @@ def f(t, y, args):
         else:
             raise
     y1 = sol.ys
-    true_y1 = jtu.tree_map(lambda x: (x * math.exp(-1))[None], y0)
+    if jnp.iscomplexobj(y_dtype):
+        true_y1 = jtu.tree_map(lambda x: (x * jnp.exp(-1j))[None], y0)
+    else:
+        true_y1 = jtu.tree_map(lambda x: (x * math.exp(-1))[None], y0)
     assert shaped_allclose(y1, true_y1, atol=1e-2, rtol=1e-2)