test adjoint

Author: lockwo

.gitignore

@@ -9,3 +9,4 @@ site/
 .idea/
 .venv/
 uv.lock
+.vscode/

diffrax/_adjoint.py

@@ -22,7 +22,8 @@
     AbstractSRK,
     AbstractStratonovichSolver,
 )
-from ._term import AbstractTerm, AdjointTerm
+from ._term import AbstractTerm, AdjointTerm, MultiTerm
+from ._typing import get_origin_no_specials
 
 
 ω = cast(Callable, ω)
@@ -847,6 +848,12 @@ def loop(
                 "`diffrax.BacksolveAdjoint` is only compatible with solvers that take "
                 "a single term."
             )
+        if get_origin_no_specials(solver.term_structure, "term_structure") is MultiTerm:
+            raise NotImplementedError(
+                f"`diffrax.BacksolveAdjoint` is not compatible with solver "
+                f"`{type(solver).__name__}`, or any solver with requirements on "
+                "the noise. For SDEs, use a basic solver such as `diffrax.Euler`."
+            )
         if event is not None:
             raise NotImplementedError(
                 "`diffrax.BacksolveAdjoint` is not compatible with events."

test/test_adjoint.py

@@ -414,6 +414,34 @@ def run(y0__args, adjoint):
     assert tree_allclose(grads1, grads3, rtol=1e-3, atol=1e-3)
 
 
+@pytest.mark.parametrize(
+    "solver",
+    (diffrax.ShARK(), diffrax.SEA(), diffrax.SRA1(), diffrax.SlowRK()),
+)
+def test_backsolve_multiterm_solver_error(solver, getkey):
+    # https://github.com/patrick-kidger/diffrax/issues/558
+    t0, t1, dt0 = 0, 1, 0.01
+    bm = diffrax.VirtualBrownianTree(
+        t0, t1, 1e-3, (2,), key=getkey(), levy_area=diffrax.SpaceTimeLevyArea
+    )
+    drift = diffrax.ODETerm(lambda t, y, args: -y)
+    diffusion = diffrax.ControlTerm(
+        lambda t, y, args: lx.DiagonalLinearOperator(0.1 * jnp.zeros_like(y)), bm
+    )
+    terms = diffrax.MultiTerm(drift, diffusion)
+
+    @eqx.filter_jit
+    @jax.grad
+    def run(y0):
+        sol = diffrax.diffeqsolve(
+            terms, solver, t0, t1, dt0, y0, adjoint=diffrax.BacksolveAdjoint()
+        )
+        return jnp.sum(cast(Array, sol.ys))
+
+    with pytest.raises(NotImplementedError, match="not compatible with solver"):
+        run(jnp.array([1.0, 2.0]))
+
+
 def test_implicit_runge_kutta_direct_adjoint():
     diffrax.diffeqsolve(
         diffrax.ODETerm(lambda t, y, args: -y),