Merge pull request #3423 from jsbrittain/jax_gpu

JaxSolver fails when using GPU support with no input parameters

Author: Saransh-cpp

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

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,
     "pygment_light_style": "xcode",
     "pygment_dark_style": "monokai",

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):