From 12d672e71f0e5b9f0d9609d3ccab8b604bd5aab4 Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Fri, 21 Mar 2025 07:38:35 +0100
Subject: [PATCH] Fixed incompatibility between step_ts/jump_ts and implicit
 solvers.

---
 diffrax/_step_size_controller/clip.py     | 35 +++++++++++++++++------
 test/test_adaptive_stepsize_controller.py | 24 ++++++++++++++++
 2 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/diffrax/_step_size_controller/clip.py b/diffrax/_step_size_controller/clip.py
index 0167e789..68264bdd 100644
--- a/diffrax/_step_size_controller/clip.py
+++ b/diffrax/_step_size_controller/clip.py
@@ -1,4 +1,4 @@
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from typing import cast, Generic, Optional, TypeVar
 
 import equinox as eqx
@@ -19,7 +19,7 @@
 from .._misc import upcast_or_raise
 from .._solution import is_okay, RESULTS
 from .._term import AbstractTerm
-from .base import AbstractStepSizeController
+from .base import AbstractAdaptiveStepSizeController
 
 
 _ControllerState = TypeVar("_ControllerState")
@@ -118,7 +118,7 @@ def cond_down(_i):
 
 
 class ClipStepSizeController(
-    AbstractStepSizeController[_ClipState[_ControllerState], _Dt0]
+    AbstractAdaptiveStepSizeController[_ClipState[_ControllerState], _Dt0]
 ):
     """Wraps an existing step controller with three pieces of functionality:
 
@@ -166,20 +166,32 @@ class ClipStepSizeController(
         ```
     """
 
-    controller: AbstractStepSizeController[_ControllerState, _Dt0]
+    controller: AbstractAdaptiveStepSizeController[_ControllerState, _Dt0]
     step_ts: Optional[Real[Array, " steps"]]
     jump_ts: Optional[Real[Array, " jumps"]]
     store_rejected_steps: Optional[int] = eqx.field(static=True)
     callback_on_reject: Optional[Callable] = eqx.field(static=True)
 
+    @property
+    def atol(self):
+        return self.controller.atol
+
+    @property
+    def rtol(self):
+        return self.controller.rtol
+
+    @property
+    def norm(self):  # pyright: ignore[reportIncompatibleMethodOverride]
+        return self.controller.norm
+
     @eqxi.doc_remove_args("_callback_on_reject")
     def __init__(
         self,
-        controller,
-        step_ts=None,
-        jump_ts=None,
-        store_rejected_steps=None,
-        _callback_on_reject=None,
+        controller: AbstractAdaptiveStepSizeController[_ControllerState, _Dt0],
+        step_ts: None | Sequence[RealScalarLike] | Real[Array, " steps"] = None,
+        jump_ts: None | Sequence[RealScalarLike] | Real[Array, " jumps"] = None,
+        store_rejected_steps: Optional[int] = None,
+        _callback_on_reject: Optional[Callable] = None,
     ):
         """**Arguments**:
 
@@ -198,6 +210,11 @@ def __init__(
             that this is not the total number of rejected steps in a solve, but just the
             maximum number of *consecutive* rejected steps.)
         """
+        if not isinstance(controller, AbstractAdaptiveStepSizeController):
+            raise ValueError(
+                "Can only apply `ClipStepSizeController` to adaptive step size "
+                f"controllers, but got {controller}."
+            )
         self.controller = controller
         self.step_ts = _none_or_sorted_array(step_ts)
         self.jump_ts = _none_or_sorted_array(jump_ts)
diff --git a/test/test_adaptive_stepsize_controller.py b/test/test_adaptive_stepsize_controller.py
index 68508a2e..8161b9d8 100644
--- a/test/test_adaptive_stepsize_controller.py
+++ b/test/test_adaptive_stepsize_controller.py
@@ -312,3 +312,27 @@ def test_find_idx_with_hint():
         assert idx == 3  # not 2; we want the first value *strictly* greater.
         idx = _find_idx_with_hint(1.9, ts, hint)
         assert idx == 2
+
+
+# https://github.com/patrick-kidger/diffrax/issues/607
+@pytest.mark.parametrize("new", (False, True))
+def test_implicit_solver_with_clip_controller(new: bool):
+    term = diffrax.ODETerm(lambda t, y, args: -y)
+    solver = diffrax.Kvaerno3()
+    if new:
+        ssc = diffrax.PIDController(rtol=1e-3, atol=1e-3)
+        ssc = diffrax.ClipStepSizeController(ssc, jump_ts=[0.5])
+    else:
+        ssc = diffrax.PIDController(jump_ts=[0.5], rtol=1e-3, atol=1e-3)  # pyright: ignore[reportCallIssue]
+    diffrax.diffeqsolve(
+        term,
+        solver,
+        t0=0,
+        t1=1,
+        dt0=0.01,
+        args=None,
+        y0=1.0,
+        stepsize_controller=ssc,
+        max_steps=16384,
+        saveat=diffrax.SaveAt(t1=True),
+    )