Fixed where a nonbatchable check was being called.

patrick-kidger · Dec 9, 2024 · 825e4e0 · 825e4e0
1 parent 1ae1d58
commit 825e4e0
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/diffrax/_solver/runge_kutta.py b/diffrax/_solver/runge_kutta.py
@@ -964,7 +964,7 @@ def eval_k_jac():
                         assert implicit_tableau.a_diagonal[0] == 0  # pyright: ignore
                         assert len(set(implicit_tableau.a_diagonal[1:])) == 1  # pyright: ignore
                         jac_stage_index = 1
-                        stage_index = eqxi.nonbatchable(stage_index)
+                    stage_index = eqxi.nonbatchable(stage_index)
                     # These `stop_gradients` are needed to work around the lack of
                     # symbolic zeros in `custom_vjp`s.
                     if eval_fs:

diff --git a/test/test_global_interpolation.py b/test/test_global_interpolation.py
@@ -340,7 +340,7 @@ def _test_dense_interpolation(solver, key, t1):
 
 
 @pytest.mark.parametrize("solver", all_ode_solvers + all_split_solvers)
-def test_dense_interpolation(solver, getkey):
+def test_dense_interpolation(solver):
     solver = implicit_tol(solver)
     key = jr.PRNGKey(5678)
     vals, true_vals, derivs, true_derivs = _test_dense_interpolation(solver, key, 1)