return RESULTS from reversible backward_step

Author: sammccallum

diffrax/_adjoint.py

@@ -1013,7 +1013,7 @@ def forward_step(y0, solver_state, args, terms):
         t1 = ts[ts_index]
         t0 = ts[ts_index - 1]
 
-        y0, dense_info, solver_state = solver.backward_step(
+        y0, dense_info, solver_state, result = solver.backward_step(
             terms, t0, t1, y1, args, solver_state, False
         )
 

diffrax/_solver/base.py

@@ -371,7 +371,7 @@ def backward_step(
         args: Args,
         solver_state: _SolverState,
         made_jump: BoolScalarLike,
-    ) -> tuple[Y, DenseInfo, _SolverState]:
+    ) -> tuple[Y, DenseInfo, _SolverState, RESULTS]:
         """
         Make a single backward step with the reversible solver.
 

diffrax/_solver/leapfrog_midpoint.py

@@ -52,9 +52,9 @@ class LeapfrogMidpoint(AbstractReversibleSolver):
     """
 
     term_structure: ClassVar = AbstractTerm
-    interpolation_cls: ClassVar[
-        Callable[..., LocalLinearInterpolation]
-    ] = LocalLinearInterpolation
+    interpolation_cls: ClassVar[Callable[..., LocalLinearInterpolation]] = (
+        LocalLinearInterpolation
+    )
 
     def order(self, terms):
         return 2
@@ -101,7 +101,7 @@ def backward_step(
         args: Args,
         solver_state: _SolverState,
         made_jump: BoolScalarLike,
-    ) -> tuple[Y, DenseInfo, _SolverState]:
+    ) -> tuple[Y, DenseInfo, _SolverState, RESULTS]:
         del made_jump
         t0, y0, dt = solver_state
         tm1 = t0 - dt
@@ -114,7 +114,7 @@ def backward_step(
         solver_state = jax.lax.cond(
             tm1 > 0, lambda _: (tm1, ym1, dt), lambda _: (t0, y0, dt), None
         )
-        return y0, dense_info, solver_state
+        return y0, dense_info, solver_state, RESULTS.successful
 
     def func(self, terms: AbstractTerm, t0: RealScalarLike, y0: Y, args: Args) -> VF:
         return terms.vf(t0, y0, args)

diffrax/_solver/reversible.py

@@ -153,19 +153,21 @@ def backward_step(
         args: Args,
         solver_state: _SolverState,
         made_jump: BoolScalarLike,
-    ) -> tuple[Y, DenseInfo, _SolverState]:
+    ) -> tuple[Y, DenseInfo, _SolverState, RESULTS]:
         original_solver_state, z1 = solver_state
-        step_y1, _, _, original_solver_state, _ = self.solver.step(
+        step_y1, _, _, original_solver_state, result1 = self.solver.step(
             terms, t1, t0, y1, args, original_solver_state, True
         )
         z0 = (ω(z1) - ω(y1) + ω(step_y1)).ω
-        step_z0, _, dense_info, _, _ = self.solver.step(
+        step_z0, _, dense_info, _, result2 = self.solver.step(
             terms, t0, t1, z0, args, original_solver_state, True
         )
         y0 = ((1 / self.coupling_parameter) * (ω(y1) - ω(step_z0)) + ω(z0)).ω
+
         solver_state = (original_solver_state, z0)
+        result = update_result(result1, result2)
 
-        return y0, dense_info, solver_state
+        return y0, dense_info, solver_state, result
 
     def func(
         self, terms: PyTree[AbstractTerm], t0: RealScalarLike, y0: Y, args: Args

diffrax/_solver/reversible_heun.py

@@ -49,9 +49,9 @@ class ReversibleHeun(
     """
 
     term_structure: ClassVar = AbstractTerm
-    interpolation_cls: ClassVar[
-        Callable[..., LocalLinearInterpolation]
-    ] = LocalLinearInterpolation  # TODO use something better than this?
+    interpolation_cls: ClassVar[Callable[..., LocalLinearInterpolation]] = (
+        LocalLinearInterpolation  # TODO use something better than this?
+    )
 
     def order(self, terms):
         return 2
@@ -104,7 +104,7 @@ def backward_step(
         args: Args,
         solver_state: _SolverState,
         made_jump: BoolScalarLike,
-    ) -> tuple[Y, DenseInfo, _SolverState]:
+    ) -> tuple[Y, DenseInfo, _SolverState, RESULTS]:
         yhat1, vf1 = solver_state
 
         control = terms.contr(t0, t1)
@@ -114,7 +114,7 @@ def backward_step(
 
         dense_info = dict(y0=y0, y1=y1)
         solver_state = (yhat0, vf0)
-        return y0, dense_info, solver_state
+        return y0, dense_info, solver_state, RESULTS.successful
 
     def func(self, terms: AbstractTerm, t0: RealScalarLike, y0: Y, args: Args) -> VF:
         return terms.vf(t0, y0, args)

diffrax/_solver/semi_implicit_euler.py

@@ -33,9 +33,9 @@ class SemiImplicitEuler(AbstractReversibleSolver):
     """
 
     term_structure: ClassVar = (AbstractTerm, AbstractTerm)
-    interpolation_cls: ClassVar[
-        Callable[..., LocalLinearInterpolation]
-    ] = LocalLinearInterpolation
+    interpolation_cls: ClassVar[Callable[..., LocalLinearInterpolation]] = (
+        LocalLinearInterpolation
+    )
 
     def order(self, terms):
         return 1
@@ -83,7 +83,7 @@ def backward_step(
         args: Args,
         solver_state: _SolverState,
         made_jump: BoolScalarLike,
-    ) -> tuple[tuple[Ya, Yb], DenseInfo, _SolverState]:
+    ) -> tuple[tuple[Ya, Yb], DenseInfo, _SolverState, RESULTS]:
         del solver_state, made_jump
 
         term_1, term_2 = terms
@@ -96,7 +96,7 @@ def backward_step(
 
         y0 = (y0_1, y0_2)
         dense_info = dict(y0=y0, y1=y1)
-        return y0, dense_info, None
+        return y0, dense_info, None, RESULTS.successful
 
     def func(
         self,