Enable implicit solvers for complex inputs (#411)

* Enable implicit solvers for complex inputs

* change version

* make pyright happy

Author: Randl

diffrax/_integrate.py

@@ -709,11 +709,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.",

diffrax/_progress_meter.py

@@ -294,7 +294,7 @@ def _step(_progress, _idx):
                 # Return the idx to thread the callbacks in the correct order.
                 return _idx
 
-        return jax.pure_callback(_step, idx, progress, idx, vectorized=True)  # pyright: ignore
+        return jax.pure_callback(_step, idx, progress, idx, vectorized=True)
 
     def close(self, close_bar: Callable[[Any], None], idx: IntScalarLike):
         def _close(_idx):

pyproject.toml

@@ -23,7 +23,7 @@ classifiers = [
     "Topic :: Scientific/Engineering :: Mathematics",
 ]
 urls = {repository = "https://github.com/patrick-kidger/diffrax" }
-dependencies = ["jax>=0.4.23", "jaxtyping>=0.2.24", "typing_extensions>=4.5.0", "typeguard==2.13.3", "equinox>=0.11.2", "lineax>=0.0.5", "optimistix>=0.0.6"]
+dependencies = ["jax>=0.4.23", "jaxtyping>=0.2.24", "typing_extensions>=4.5.0", "typeguard==2.13.3", "equinox>=0.11.2", "lineax>=0.0.5", "optimistix>=0.0.7"]
 
 [build-system]
 requires = ["hatchling"]

test/test_integrate.py

@@ -77,13 +77,10 @@ def test_basic(solver, t_dtype, y_dtype, treedef, stepsize_controller, getkey):
         return
 
     if jnp.iscomplexobj(y_dtype) and treedef != jtu.tree_structure(None):
-        if isinstance(solver, diffrax.AbstractImplicitSolver):
-            return
-        else:
-            complex_warn = pytest.warns(match="Complex dtype")
+        complex_warn = pytest.warns(match="Complex dtype")
 
-            def f(t, y, args):
-                return jtu.tree_map(lambda yi: -1j * yi, y)
+        def f(t, y, args):
+            return jtu.tree_map(lambda yi: -1j * yi, y)
     else:
         complex_warn = contextlib.nullcontext()
 
@@ -152,8 +149,6 @@ def test_ode_order(solver, dtype):
 
     A = jr.normal(akey, (10, 10), dtype=dtype) * 0.5
 
-    if jnp.iscomplexobj(A) and isinstance(solver, diffrax.AbstractImplicitSolver):
-        return
     if (
         solver.term_structure
         == diffrax.MultiTerm[tuple[diffrax.AbstractTerm, diffrax.AbstractTerm]]

test/test_interpolation.py

@@ -60,8 +60,6 @@ def test_derivative(dtype, getkey):
         solver = implicit_tol(solver)
         y0 = jr.normal(getkey(), (3,), dtype=dtype)
 
-        if jnp.iscomplexobj(y0) and isinstance(solver, diffrax.AbstractImplicitSolver):
-            continue
         solution = diffrax.diffeqsolve(
             diffrax.ODETerm(lambda t, y, p: -y),
             solver,
@@ -77,8 +75,7 @@ def test_derivative(dtype, getkey):
     for solver in all_split_solvers:
         solver = implicit_tol(solver)
         y0 = jr.normal(getkey(), (3,), dtype=dtype)
-        if jnp.iscomplexobj(y0) and isinstance(solver, diffrax.AbstractImplicitSolver):
-            continue
+
         solution = diffrax.diffeqsolve(
             diffrax.MultiTerm(
                 diffrax.ODETerm(lambda t, y, p: -0.7 * y),