From 013c8a43a87f4ccb2c1fb3d567ba12fd1c0e16ec Mon Sep 17 00:00:00 2001
From: AHMETHAKANBEZIR1 <ahmethakan0008@gmail.com>
Date: Fri, 9 Oct 2026 14:57:36 +0300
Subject: [PATCH] Fix rounded rejected endpoints in PIDController

---
 diffrax/_step_size_controller/pid.py      |  8 +++++
 test/test_adaptive_stepsize_controller.py | 40 +++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/diffrax/_step_size_controller/pid.py b/diffrax/_step_size_controller/pid.py
index 752b9bed..442d5f70 100644
--- a/diffrax/_step_size_controller/pid.py
+++ b/diffrax/_step_size_controller/pid.py
@@ -556,6 +556,14 @@ def _scale(_y0, _y1_candidate, _y_error):
 
         next_t0 = jnp.where(keep_step, t1, t0)
         next_t1 = next_t0 + dt
+        if self.dtmin is None:
+            # A smaller dt can round back to the same rejected endpoint. Move to
+            # the previous representable time if it still gives a nonzero step.
+            smaller_t1 = eqxi.prevbefore(
+                jnp.asarray(t1, dtype=jnp.result_type(next_t1))
+            )
+            rounded_rejection = (~keep_step) & (next_t1 >= t1) & (smaller_t1 > next_t0)
+            next_t1 = jnp.where(rounded_rejection, smaller_t1, next_t1)
 
         inv_scaled_error = jnp.where(keep_step, inv_scaled_error, prev_inv_scaled_error)
         prev_inv_scaled_error = jnp.where(
diff --git a/test/test_adaptive_stepsize_controller.py b/test/test_adaptive_stepsize_controller.py
index 4a51336f..03936f58 100644
--- a/test/test_adaptive_stepsize_controller.py
+++ b/test/test_adaptive_stepsize_controller.py
@@ -446,3 +446,43 @@ def test_t0_at_jump_time():
     # What used to happen was something very weird where we'd oscillate across the
     # jump time.
     assert sol.result == diffrax.RESULTS.successful
+
+
+@pytest.mark.parametrize("time_sign", [-1, 1])
+def test_rejected_step_rounding(time_sign):
+    # Four ULPs leave room to reduce the step, even when the PID factor rounds
+    # the proposed endpoint back to the endpoint that was just rejected.
+    t0 = jnp.asarray(time_sign * 100_000, dtype=jnp.float32)
+    t1 = t0 + jnp.float32(0.03125)
+    controller = diffrax.PIDController(rtol=0, atol=1)
+    state = (jnp.float32(1), jnp.float32(1), False)
+    keep, next_t0, next_t1, _, _, result = eqx.filter_jit(controller.adapt_step_size)(
+        t0, t1, jnp.float32(0), jnp.float32(0), None, jnp.float32(1.1), 5, state
+    )
+    assert not keep
+    assert next_t0 == t0
+    assert t0 < next_t1 < t1
+    assert result == diffrax.RESULTS.successful
+
+
+@pytest.mark.parametrize("backwards", [False, True])
+def test_large_time_rejected_step(backwards):
+    # Translating an autonomous ODE should not leave the solver repeatedly
+    # rejecting exactly the same representable interval.
+    t0 = jnp.float32(100_000)
+    sign = -1 if backwards else 1
+    term = diffrax.ODETerm(lambda t, y, args: -sign * 20 * y)
+    solution = diffrax.diffeqsolve(
+        term,
+        diffrax.Dopri5(),
+        t0=t0,
+        t1=t0 + sign * jnp.float32(0.125),
+        dt0=sign * jnp.float32(0.03125),
+        y0=jnp.float32(1),
+        stepsize_controller=diffrax.PIDController(rtol=0, atol=6e-5),
+        max_steps=64,
+        throw=False,
+    )
+    assert solution.result == diffrax.RESULTS.successful
+    assert solution.stats["num_rejected_steps"] > 0
+    assert jnp.allclose(cast(Array, solution.ys), jnp.exp(-2.5), atol=6e-5, rtol=0)