Merge pull request #3380 from BradyPlanden/jax-metal

Jax on Metal w/ GPU support
pybamm-team · Oct 2, 2023 · f81de94 · f81de94
2 parents 7eb5879 + 577f9a1
commit f81de94
Show file tree

Hide file tree

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

diff --git a/pybamm/solvers/jax_bdf_solver.py b/pybamm/solvers/jax_bdf_solver.py
@@ -18,7 +18,9 @@
     from jax.tree_util import tree_flatten, tree_map, tree_unflatten
     from jax.util import cache, safe_map, split_list
 
-    config.update("jax_enable_x64", True)
+    platform = jax.lib.xla_bridge.get_backend().platform.casefold()
+    if platform != "metal":
+        config.update("jax_enable_x64", True)
 
     MAX_ORDER = 5
     NEWTON_MAXITER = 4

diff --git a/pybamm/solvers/jax_solver.py b/pybamm/solvers/jax_solver.py
@@ -227,7 +227,11 @@ async def solve_model_async(inputs_v):
                 return await asyncio.gather(*coro)
 
             y = asyncio.run(solve_model_for_inputs())
-        elif platform.startswith("gpu") or platform.startswith("tpu"):
+        elif (
+            platform.startswith("gpu")
+            or platform.startswith("tpu")
+            or platform.startswith("metal")
+        ):
             # gpu execution runs faster when parallelised with vmap
             # (see also comment below regarding single-program multiple-data
             #  execution (SPMD) using pmap on multiple XLAs)