Fix JAX deprecation warning

Author: patrick-kidger

diffrax/_global_interpolation.py

@@ -40,7 +40,7 @@ def _interpret_t(
     ) -> tuple[IntScalarLike, RealScalarLike]:
         maxlen = self.ts_size - 2
         index = jnp.searchsorted(self.ts, t, side="left" if left else "right")
-        index = jnp.clip(index - 1, a_min=0, a_max=maxlen)
+        index = jnp.clip(index - 1, min=0, max=maxlen)
         # Will never access the final element of `ts`; this is correct behaviour.
         fractional_part = t - self.ts[index]
         return index, fractional_part

diffrax/_step_size_controller/adaptive.py

@@ -602,8 +602,8 @@ def _scale(_y0, _y1_candidate, _y_error):
         factormin = jnp.where(keep_step, 1, self.factormin)
         factor = jnp.clip(
             self.safety * factor1 * factor2 * factor3,
-            a_min=factormin,
-            a_max=self.factormax,
+            min=factormin,
+            max=self.factormax,
         )
         # Once again, see above. In case we have gradients on {i,p,d}coeff.
         # (Probably quite common for them to have zero tangents if passed across

pyproject.toml

@@ -23,7 +23,7 @@ classifiers = [
     "Topic :: Scientific/Engineering :: Mathematics",
 ]
 urls = {repository = "https://github.com/patrick-kidger/diffrax" }
-dependencies = ["jax>=0.4.23", "jaxtyping>=0.2.24", "typing_extensions>=4.5.0", "typeguard==2.13.3", "equinox>=0.11.2", "lineax>=0.0.5", "optimistix>=0.0.7"]
+dependencies = ["jax>=0.4.28", "jaxtyping>=0.2.24", "typing_extensions>=4.5.0", "typeguard==2.13.3", "equinox>=0.11.2", "lineax>=0.0.5", "optimistix>=0.0.7"]
 
 [build-system]
 requires = ["hatchling"]