patched edge-case in LinearInterpolation

Author: packquickly

diffrax/global_interpolation.py

@@ -125,10 +125,13 @@ def _index(_ys):
         prev_t = self.ts[index]
         next_t = self.ts[index + 1]
         diff_t = next_t - prev_t
-
-        return (
-            prev_ys**ω + (next_ys**ω - prev_ys**ω) * (fractional_part / diff_t)
-        ).ω
+        diff_nonzero = diff_t >= jnp.finfo(diff_t.dtype).eps
+        safe_diff = jnp.where(diff_nonzero, diff_t, jnp.ones_like(diff_t))
+        return jnp.where(
+            diff_nonzero,
+            (prev_ys**ω + (next_ys**ω - prev_ys**ω) * (fractional_part / safe_diff)).ω,
+            prev_ys
+        )
 
     @eqx.filter_jit
     def derivative(self, t: Scalar, left: bool = True) -> PyTree: