Merge pull request #3671 from prady0t/fix-jax

Fixed Deprecation warnings of `jax_bdf_solver`
pybamm-team · Jan 3, 2024 · ad3b816 · ad3b816
2 parents 1c1bcdd + 7110cac
commit ad3b816
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pybamm/solvers/jax_bdf_solver.py b/pybamm/solvers/jax_bdf_solver.py
@@ -10,7 +10,7 @@
     import jax
     import jax.numpy as jnp
     from jax import core, dtypes
-    from jax import linear_util as lu
+    from jax.extend import linear_util as lu
     from jax.api_util import flatten_fun_nokwargs
     from jax.config import config
     from jax.flatten_util import ravel_pytree