restrict Reversible to AbstractERK and check result in adjoint

Author: sammccallum

diffrax/_adjoint.py

@@ -16,6 +16,7 @@
 
 from ._heuristics import is_sde, is_unsafe_sde
 from ._saveat import save_y, SaveAt, SubSaveAt
+from ._solution import RESULTS
 from ._solver import (
     AbstractItoSolver,
     AbstractRungeKutta,
@@ -994,9 +995,10 @@ def _loop_reversible_bwd(
 
     def grad_step(state):
         def forward_step(y0, solver_state, args, terms):
-            y1, _, dense_info, new_solver_state, _ = solver.step(
+            y1, _, dense_info, new_solver_state, result = solver.step(
                 terms, t0, t1, y0, args, solver_state, False
             )
+            assert result == RESULTS.successful
             return y1, dense_info, new_solver_state
 
         (
@@ -1016,6 +1018,7 @@ def forward_step(y0, solver_state, args, terms):
         y0, dense_info, solver_state, result = solver.backward_step(
             terms, t0, t1, y1, args, solver_state, False
         )
+        assert result == RESULTS.successful
 
         # Pull gradients back through interpolation
 

diffrax/_solver/reversible.py

@@ -8,10 +8,10 @@
 from .._solution import RESULTS, update_result
 from .._solver.base import (
     AbstractReversibleSolver,
-    AbstractSolver,
     AbstractWrappedSolver,
 )
 from .._term import AbstractTerm
+from .runge_kutta import AbstractERK
 
 
 ω = cast(Callable, ω)
@@ -26,7 +26,7 @@ class Reversible(
     """
     Reversible solver method.
 
-    Allows any solver ([`diffrax.AbstractSolver`][]) to be made
+    Allows any explicit Runge-Kutta solver ([`diffrax.AbstractERK`][]) to be made
     algebraically reversible.
 
     **Arguments:**
@@ -77,7 +77,7 @@ class Reversible(
         ```
     """
 
-    solver: AbstractSolver
+    solver: AbstractERK
     coupling_parameter: float = 0.999
 
     @property
@@ -114,6 +114,10 @@ def init(
         y0: Y,
         args: Args,
     ) -> _SolverState:
+        if not isinstance(self.solver, AbstractERK):
+            raise ValueError(
+                "`Reversible` is only compatible with `AbstractERK` base solvers."
+            )
         original_solver_init = self.solver.init(terms, t0, t1, y0, args)
         return (original_solver_init, y0)
 

test/test_reversible.py

@@ -198,34 +198,6 @@ def test_reversible_explicit(stepsize_controller, saveat):
     )
 
 
-@pytest.mark.parametrize(
-    "saveat",
-    [
-        diffrax.SaveAt(t0=True, t1=True),
-        diffrax.SaveAt(t0=True, ts=jnp.linspace(0, 5, 10), t1=True),
-    ],
-)
-def test_reversible_implicit(saveat):
-    n = 10
-    y0 = jnp.linspace(1, 10, num=n)
-    key = jr.PRNGKey(10)
-    f = VectorField(n, n, n, depth=4, key=key)
-    terms = diffrax.ODETerm(f)
-    args = jnp.array([0.5])
-    base_solver = diffrax.Kvaerno5()
-    solver = diffrax.Reversible(base_solver)
-    stepsize_controller = diffrax.PIDController(rtol=1e-8, atol=1e-8)
-
-    _compare_grads(
-        (y0, args, terms),
-        solver,
-        solver,
-        saveat,
-        stepsize_controller,
-        dual_y0=False,
-    )
-
-
 @pytest.mark.parametrize(
     "saveat",
     [