Fixes for JAX 0.4.36 which changes the name of an error.

Author: patrick-kidger

test/test_adjoint.py

@@ -4,7 +4,7 @@
 import diffrax
 import equinox as eqx
 import jax
-import jax.interpreters.ad
+import jax._src.interpreters.ad
 import jax.numpy as jnp
 import jax.random as jr
 import jax.tree_util as jtu
@@ -21,6 +21,7 @@ class _VectorField(eqx.Module):
     diff_arg: float
 
     def __call__(self, t, y, args):
+        del t
         assert y.shape == (2,)
         diff_arg, nondiff_arg = args
         dya = diff_arg * y[0] + nondiff_arg * y[1]
@@ -29,7 +30,7 @@ def __call__(self, t, y, args):
 
 
 @pytest.mark.slow
-def test_against(getkey):
+def test_against():
     y0 = jnp.array([0.9, 5.4])
     args = (0.1, -1)
     term = diffrax.ODETerm(_VectorField(nondiff_arg=1, diff_arg=-0.1))
@@ -215,6 +216,7 @@ def test_closure_errors():
     @eqx.filter_value_and_grad
     def run(model):
         def f(t, y, args):
+            del t, args
             return model(y)
 
         sol = diffrax.diffeqsolve(
@@ -228,7 +230,7 @@ def f(t, y, args):
         )
         return jnp.sum(cast(Array, sol.ys))
 
-    with pytest.raises(jax.interpreters.ad.CustomVJPException):
+    with pytest.raises(jax._src.interpreters.ad.CustomVJPException):
         run(mlp)
 
 
@@ -239,6 +241,7 @@ class VectorField(eqx.Module):
         model: Callable
 
         def __call__(self, t, y, args):
+            del t, args
             return self.model(y)
 
     @eqx.filter_jit
@@ -307,12 +310,12 @@ def make_step(model, opt_state, target_steady_state):
         model = eqx.apply_updates(model, updates)
         return model, opt_state
 
-    for step in range(100):
+    for _ in range(100):
         model, opt_state = make_step(model, opt_state, target_steady_state)
     assert tree_allclose(model.steady_state, target_steady_state, rtol=1e-2, atol=1e-2)
 
 
-def test_backprop_ts(getkey):
+def test_backprop_ts():
     mlp = eqx.nn.MLP(1, 1, 8, 2, key=jr.PRNGKey(0))
 
     @eqx.filter_jit
@@ -338,14 +341,17 @@ def run(model):
 )
 def test_sde_against(diffusion_fn, getkey):
     def f(t, y, args):
+        del t
         k0, _ = args
         return -k0 * y
 
     def g(t, y, args):
+        del t
         _, k1 = args
         return k1 * y
 
     def g_lx(t, y, args):
+        del t
         _, k1 = args
         return lx.DiagonalLinearOperator(k1 * y)