Merge pull request #3423 from jsbrittain/jax_gpu

JaxSolver fails when using GPU support with no input parameters
pybamm-team · Oct 31, 2023 · 138cbf2 · 138cbf2
2 parents 618b481 + 6cc3940
commit 138cbf2
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # [Unreleased](https://github.com/pybamm-team/PyBaMM/)
 
+## Bug fixes
+
+- Fixed a bug where the JaxSolver would fails when using GPU support with no input parameters ([#3423](https://github.com/pybamm-team/PyBaMM/pull/3423))
+
 # [v23.9rc0](https://github.com/pybamm-team/PyBaMM/tree/v23.9rc0) - 2023-10-31
 
 ## Features

diff --git a/docs/conf.py b/docs/conf.py
@@ -154,6 +154,7 @@
     "navbar_end": ["theme-switcher", "navbar-icon-links"],
     # add Algolia to the persistent navbar, this removes the default search icon
     "navbar_persistent": "algolia-searchbox",
+    "navigation_with_keys": False,
     "use_edit_page_button": True,
     "analytics": {
         "plausible_analytics_domain": "docs.pybamm.org",

diff --git a/pybamm/solvers/jax_solver.py b/pybamm/solvers/jax_solver.py
@@ -215,7 +215,7 @@ def _integrate(self, model, t_eval, inputs=None):
 
         y = []
         platform = jax.lib.xla_bridge.get_backend().platform.casefold()
-        if platform.startswith("cpu"):
+        if len(inputs) <= 1 or platform.startswith("cpu"):
             # cpu execution runs faster when multithreaded
             async def solve_model_for_inputs():
                 async def solve_model_async(inputs_v):