diff --git a/.gitignore b/.gitignore index 176aa2af..f35eac8a 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,4 @@ site/ .idea/ .venv/ uv.lock +.vscode/ \ No newline at end of file diff --git a/diffrax/_adjoint.py b/diffrax/_adjoint.py index 7bc081b9..fea51e5d 100644 --- a/diffrax/_adjoint.py +++ b/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." diff --git a/test/test_adjoint.py b/test/test_adjoint.py index 9e17e535..fd5baaca 100644 --- a/test/test_adjoint.py +++ b/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),