Merge pull request #3684 from cringeyburger/issue-3683-update-jax-imp…

…orts Update JAX Imports
pybamm-team · Jan 3, 2024 · b2e852e · b2e852e
2 parents 3bd05c4 + ca74358
commit b2e852e
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/pybamm/expression_tree/operations/evaluate_python.py b/pybamm/expression_tree/operations/evaluate_python.py
@@ -11,11 +11,10 @@
 
 if pybamm.have_jax():
     import jax
-    from jax.config import config
 
     platform = jax.lib.xla_bridge.get_backend().platform.casefold()
     if platform != "metal":
-        config.update("jax_enable_x64", True)
+        jax.config.update("jax_enable_x64", True)
 
 
 class JaxCooMatrix:

diff --git a/pybamm/solvers/jax_bdf_solver.py b/pybamm/solvers/jax_bdf_solver.py
@@ -12,15 +12,14 @@
     from jax import core, dtypes
     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
     from jax.interpreters import partial_eval as pe
     from jax.tree_util import tree_flatten, tree_map, tree_unflatten
     from jax.util import cache, safe_map, split_list
 
     platform = jax.lib.xla_bridge.get_backend().platform.casefold()
     if platform != "metal":
-        config.update("jax_enable_x64", True)
+        jax.config.update("jax_enable_x64", True)
 
     MAX_ORDER = 5
     NEWTON_MAXITER = 4