Support Python 3.14, drop JAX on macOS Intel (#5374)

BradyPlanden · pre-commit-ci[bot] · web-flow · commit 6d6f51a946f6 · 2026-02-06T16:36:15.000Z
* feat: support Python 3.14, bump JAX and CasADi, drop JAX on macOS Intel

Bump requires-python to &lt;3.15, CasADi to 3.7.2, and JAX to &gt;=0.7.0,&lt;0.9.0
for Python 3.14 compatibility. Add a PEP 508 platform marker to exclude
JAX on macOS x86_64 since JAX dropped Intel Mac wheels in 0.5.0. Remove
explicit casadi pin since it is managed transitively through
pybammsolvers.Update idaklu_jax.py for JAX 0.7+ API changes (explicit
np.int64 casts, type check adjustments). Add raise_jax_not_found() helper
with a platform-aware error message so macOS Intel users get a clear explanation.

* docs: add changelog entry

* fix: update deprecated JAX API calls for JAX 0.7+

Migrate jax.lib.xla_bridge.get_backend() to jax.extend.backend.get_backend()
(deprecated in JAX 0.8.0) and jax.interpreters.mlir.custom_call to
jax._src.interpreters.mlir.custom_call.

* style: pre-commit fixes

* fix: add Python &gt;=3.11 marker for JAX and handle 3.14 math error message

Add python_version &gt;= '3.11' to the JAX dependency marker since JAX 0.7+
requires Python 3.11. Update test_bruggeman_invalid_values_raise regex to
match both the old 'math domain error' and the new Python 3.14 message
'expected a positive input' (cpython#101410).

* fix: skip jax parametrized test when jax is not installed

* infra: increment pybammsolvers

* infra: add Python 3.14 to CI, docs, and packaging

Update all workflows, documentation, Dockerfile, ReadTheDocs config,
and pyproject.toml classifiers to include Python 3.14. Drop Python 3.10
from CI test matrices.

---------

Co-authored-by: pre-commit-ci[bot] &lt;66853113+pre-commit-ci[bot]@users.noreply.github.com&gt;
diff --git a/.github/workflows/periodic_benchmarks.yml b/.github/workflows/periodic_benchmarks.yml
@@ -32,10 +32,10 @@ jobs:
         with:
           persist-credentials: false
 
-      - name: Set up Python 3.12
+      - name: Set up Python 3.14
         uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
         with:
-          python-version: 3.12
+          python-version: 3.14
 
       - name: Install Linux system dependencies
         run: |
@@ -70,10 +70,10 @@ jobs:
     permissions:
       contents: write
     steps:
-      - name: Set up Python 3.12
+      - name: Set up Python 3.14
         uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
         with:
-          python-version: 3.12
+          python-version: 3.14
 
       - name: Install asv
         run: pip install asv
diff --git a/.github/workflows/publish_pypi.yml b/.github/workflows/publish_pypi.yml
@@ -18,7 +18,7 @@ jobs:
           persist-credentials: true
       - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
         with:
-          python-version: 3.13
+          python-version: 3.14
 
       - name: Build wheel
         run: pipx run build --outdir deploy
diff --git a/.github/workflows/run_benchmarks_over_history.yml b/.github/workflows/run_benchmarks_over_history.yml
@@ -40,10 +40,10 @@ jobs:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
         with:
           persist-credentials: false
-      - name: Set up Python 3.12
+      - name: Set up Python 3.14
         uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
         with:
-          python-version: 3.12
+          python-version: 3.14
 
       - name: Install nox and asv
         run: pip install -U pip nox asv
@@ -131,10 +131,10 @@ jobs:
     permissions:
       contents: write
     steps:
-      - name: Set up Python 3.12
+      - name: Set up Python 3.14
         uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
         with:
-          python-version: 3.12
+          python-version: 3.14
 
       - name: Install asv
         run: pip install asv
diff --git a/.github/workflows/run_periodic_tests.yml b/.github/workflows/run_periodic_tests.yml
@@ -31,7 +31,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ ubuntu-latest, macos-15-intel, macos-latest, windows-latest ]
-        python-version: ["3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.11", "3.12", "3.13", "3.14"]
     name: Tests (${{ matrix.os }} / Python ${{ matrix.python-version }})
 
     steps:
diff --git a/.github/workflows/test_on_push.yml b/.github/workflows/test_on_push.yml
@@ -49,7 +49,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, macos-15-intel, macos-latest, windows-latest]
-        python-version: ["3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.11", "3.12", "3.13", "3.14"]
     name: Tests (${{ matrix.os }} / Python ${{ matrix.python-version }})
 
     steps:
diff --git a/.readthedocs.yaml b/.readthedocs.yaml
@@ -19,7 +19,7 @@ build:
     - "graphviz"
   os: ubuntu-22.04
   tools:
-    python: "3.13"
+    python: "3.14"
   jobs:
     pre_create_environment:
       - asdf plugin add uv
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## Features
 
+- Added support for Python 3.14. ([#5374](https://github.com/pybamm-team/PyBaMM/pull/5374))
 - Added regularisation to the kinetics and OCPs so they are more numerically stable. ([#5371](https://github.com/pybamm-team/PyBaMM/pull/5371))
 - Improve the performance of matrix multiplication with CasADi expressions. ([#5351](https://github.com/pybamm-team/PyBaMM/pull/5351))
 - Adds option for lists of inputs to `solve` to include input parameters which are used
@@ -12,6 +13,7 @@ as initial conditions. ([#5311](https://github.com/pybamm-team/PyBaMM/pull/5311)
 
 ## Breaking changes
 
+- Dropped JAX support on macOS with Intel (x86_64) processors. JAX dropped macOS Intel wheels in version 0.5.0, and the minimum JAX version has been bumped to >=0.7.0 for Python 3.14 compatibility. macOS users require Apple Silicon (M-series) for JAX features. ([#5374](https://github.com/pybamm-team/PyBaMM/pull/5374))
 - Added a small regularisation term to the exchange current density which slightly modifies the functional form of the kinetics as stoichiometry approaches 0 or 1. ([#5371](https://github.com/pybamm-team/PyBaMM/pull/5371))
 - The default OCP barrier as stoichiometry approaches 0 or 1 is now smooth instead of asymptotically approaching infinity. This may change the behavior of ESOH solvers in extreme stoichiometry limits. ([#5371](https://github.com/pybamm-team/PyBaMM/pull/5371))
 - Migrated `docs` and `dev` dependencies from `project.optional-dependencies` to `dependency-groups` per PEP 735. ([#5368](https://github.com/pybamm-team/PyBaMM/pull/5368))
diff --git a/docs/source/user_guide/installation/gnu-linux-mac.rst b/docs/source/user_guide/installation/gnu-linux-mac.rst
@@ -6,7 +6,7 @@ GNU/Linux & macOS
 Prerequisites
 -------------
 
-To use PyBaMM, you must have Python 3.10, 3.11, 3.12, 3.13 installed.
+To use PyBaMM, you must have Python 3.10, 3.11, 3.12, 3.13, or 3.14 installed.
 
 .. tab:: Debian-based distributions (Debian, Ubuntu)
 
@@ -43,7 +43,7 @@ User install
 
 We recommend to install PyBaMM within a virtual environment, in order
 not to alter any distribution Python files.
-First, make sure you are using Python 3.10, 3.11, 3.12, 3.13.
+First, make sure you are using Python 3.10, 3.11, 3.12, 3.13, or 3.14.
 To create a virtual environment ``env`` within your current directory type:
 
 .. code:: bash
diff --git a/docs/source/user_guide/installation/index.rst b/docs/source/user_guide/installation/index.rst
@@ -213,7 +213,7 @@ Dependency                                                  Minimum Version    p
 Jax dependencies
 ^^^^^^^^^^^^^^^^
 
-Installable with ``pip install "pybamm[jax]"``, currently supported on Python 3.10-3.12.
+Installable with ``pip install "pybamm[jax]"``, currently supported on Python 3.10-3.14.
 
 ========================================================================= ================== ================== =======================
 Dependency                                                                Minimum Version    pip extra          Notes
diff --git a/docs/source/user_guide/installation/install-from-source.rst b/docs/source/user_guide/installation/install-from-source.rst
@@ -28,7 +28,7 @@ or download the source archive on the repository's homepage.
 
 To install PyBaMM, you will need:
 
-- Python 3 (PyBaMM supports versions 3.10, 3.11, 3.12, and 3.13)
+- Python 3 (PyBaMM supports versions 3.10 — 3.14)
 - The Python headers file for your current Python version.
 - A BLAS library (for instance `openblas <https://www.openblas.net/>`_).
 - A C compiler (ex: ``gcc``).
diff --git a/docs/source/user_guide/installation/windows.rst b/docs/source/user_guide/installation/windows.rst
@@ -6,7 +6,7 @@ Windows
 Prerequisites
 -------------
 
-To use PyBaMM, you must have Python 3.10, 3.11, or 3.12 installed.
+To use PyBaMM, you must have Python 3.10 — 3.14 installed.
 
 To install Python 3 download the installation files from `Python’s
 website <https://www.python.org/downloads/windows/>`__. Make sure to
diff --git a/pyproject.toml b/pyproject.toml
@@ -9,7 +9,7 @@ license = { file = "LICENSE.txt" }
 description = "Python Battery Mathematical Modelling"
 authors = [{ name = "The PyBaMM Team", email = "pybamm@pybamm.org" }]
 maintainers = [{ name = "The PyBaMM Team", email = "pybamm@pybamm.org" }]
-requires-python = ">=3.10, <3.14"
+requires-python = ">=3.10, <3.15"
 readme = { file = "README.md", content-type = "text/markdown" }
 classifiers = [
     "Development Status :: 5 - Production/Stable",
@@ -23,14 +23,14 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
     "Topic :: Scientific/Engineering",
 ]
 dependencies = [
-    "pybammsolvers>=0.3.3,<0.4.0",
+    "pybammsolvers>=0.3.3,<0.5.0",
     "black",
     "numpy",
     "scipy>=1.11.4",
-    "casadi==3.6.7",
     "xarray>=2022.6.0",
     "anytree>=2.8.0",
     "sympy>=1.12",
@@ -63,7 +63,7 @@ cite = ["pybtex>=0.25.0"]
 bpx = ["bpx>=0.5.0,<0.6.0"]
 # Low-overhead progress bars
 tqdm = ["tqdm"]
-jax = ["jax>=0.4.36,<0.7.0"]
+jax = ["jax>=0.7.0, <0.9.0; python_version >= '3.11' and (sys_platform != 'darwin' or platform_machine != 'x86_64')"]
 # Contains all optional dependencies, except for jax, and dev dependencies
 all = [
     "scikit-fem>=8.1.0",
diff --git a/scripts/Dockerfile b/scripts/Dockerfile
@@ -1,4 +1,4 @@
-FROM python:3.12-slim-bullseye
+FROM python:3.14-slim-bookworm
 COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv
 
 WORKDIR /
diff --git a/src/pybamm/__init__.py b/src/pybamm/__init__.py
@@ -11,6 +11,7 @@
 from .util import (
     get_parameters_filepath,
     has_jax,
+    raise_jax_not_found,
     import_optional_dependency,
 )
 from .logger import logger, set_logging_level, get_new_logger
diff --git a/src/pybamm/expression_tree/operations/evaluate_python.py b/src/pybamm/expression_tree/operations/evaluate_python.py
@@ -14,8 +14,9 @@
 
 if pybamm.has_jax():
     import jax
+    import jax.extend
 
-    platform = jax.lib.xla_bridge.get_backend().platform.casefold()
+    platform = jax.extend.backend.get_backend().platform.casefold()
     if platform != "metal":
         jax.config.update("jax_enable_x64", True)
 
diff --git a/src/pybamm/solvers/idaklu_jax.py b/src/pybamm/solvers/idaklu_jax.py
@@ -23,8 +23,8 @@
     except ImportError:
         from jax.extend import ffi
     from jax import numpy as jnp
+    from jax._src.interpreters.mlir import custom_call
     from jax.interpreters import ad, batching, mlir
-    from jax.interpreters.mlir import custom_call
     from jax.tree_util import tree_flatten
 
     # Handle JAX version compatibility for Primitive location
@@ -58,9 +58,7 @@ def __init__(
         t_interp=None,
     ):
         if not pybamm.has_jax():
-            raise ModuleNotFoundError(
-                "Jax or jaxlib is not installed, please see https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-jaxsolver"
-            )  # pragma: no cover
+            pybamm.raise_jax_not_found()  # pragma: no cover
         self.jaxpr = (
             None  # JAX expression representing the IDAKLU-wrapped solver object
         )
@@ -491,10 +489,10 @@ def _jax_vjp_impl(
         t = primals[0]
         inputs = primals[1:]
 
-        if isinstance(invar, float):
-            invar = round(invar)
-        if isinstance(t, float):
-            t = np.array(t)
+        if not isinstance(invar, (int, str)):
+            invar = int(invar)
+        if not isinstance(t, (np.ndarray, jnp.ndarray)):
+            t = np.asarray(t)
 
         if t.ndim == 0 or (t.ndim == 1 and t.shape[0] == 1):
             # scalar time input
@@ -727,11 +725,13 @@ def f_lowering_cpu(ctx, t, *inputs):
                 # The inputs
                 operands=[
                     mlir.ir_constant(
-                        self.idaklu_jax_obj.get_index()
+                        np.int64(self.idaklu_jax_obj.get_index())
                     ),  # solver index reference
                     mlir.ir_constant(size_t),  # 'size' argument
-                    mlir.ir_constant(len(self.jax_output_variables)),  # 'vars' argument
-                    mlir.ir_constant(len(inputs)),  # 'vars' argument
+                    mlir.ir_constant(
+                        np.int64(len(self.jax_output_variables))
+                    ),  # 'vars' argument
+                    mlir.ir_constant(np.int64(len(inputs))),  # 'vars' argument
                     t,
                     *inputs,
                 ],
@@ -930,11 +930,13 @@ def f_jvp_lowering_cpu(ctx, *args):
                 # The inputs
                 operands=[
                     mlir.ir_constant(
-                        self.idaklu_jax_obj.get_index()
+                        np.int64(self.idaklu_jax_obj.get_index())
                     ),  # solver index reference
                     mlir.ir_constant(size_t),  # 'size' argument
-                    mlir.ir_constant(len(self.jax_output_variables)),  # 'vars' argument
-                    mlir.ir_constant(len(inputs_primals)),  # 'vars' argument
+                    mlir.ir_constant(
+                        np.int64(len(self.jax_output_variables))
+                    ),  # 'vars' argument
+                    mlir.ir_constant(np.int64(len(inputs_primals))),  # 'vars' argument
                     t_primal,  # 't'
                     *inputs_primals,  # inputs
                     t_tangent,  # 't'
@@ -1063,13 +1065,17 @@ def f_vjp_lowering_cpu(ctx, y_bar, invar, *primals):
                 # The inputs
                 operands=[
                     mlir.ir_constant(
-                        self.idaklu_jax_obj.get_index()
+                        np.int64(self.idaklu_jax_obj.get_index())
                     ),  # solver index reference
                     mlir.ir_constant(size_t),  # 'size' argument
-                    mlir.ir_constant(len(self.jax_inputs)),  # number of inputs
-                    mlir.ir_constant(dims_y_bar[0]),  # 'y_bar' shape[0]
-                    mlir.ir_constant(  # 'y_bar' shape[1]
-                        dims_y_bar[1] if len(dims_y_bar) > 1 else -1
+                    mlir.ir_constant(
+                        np.int64(len(self.jax_inputs))
+                    ),  # number of inputs
+                    mlir.ir_constant(np.int64(dims_y_bar[0])),  # 'y_bar' shape[0]
+                    mlir.ir_constant(
+                        np.int64(  # 'y_bar' shape[1]
+                            dims_y_bar[1] if len(dims_y_bar) > 1 else -1
+                        )
                     ),  # 'y_bar' argument
                     y_bar,  # 'y_bar'
                     invar,  # 'invar'
diff --git a/src/pybamm/solvers/jax_bdf_solver.py b/src/pybamm/solvers/jax_bdf_solver.py
@@ -32,7 +32,9 @@ def split_list(lst, indices):
         result.append(lst[start:])
         return result
 
-    platform = jax.lib.xla_bridge.get_backend().platform.casefold()
+    import jax.extend
+
+    platform = jax.extend.backend.get_backend().platform.casefold()
     if platform != "metal":
         jax.config.update("jax_enable_x64", True)
 
@@ -977,9 +979,7 @@ def jax_bdf_integrate(func, y0, t_eval, *args, rtol=1e-6, atol=1e-6, mass=None):
 
     """
     if not pybamm.has_jax():
-        raise ModuleNotFoundError(
-            "Jax or jaxlib is not installed, please see https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-jaxsolver"
-        )
+        pybamm.raise_jax_not_found()
 
     def _check_arg(arg):
         if not isinstance(arg, core.Tracer) and not core.valid_jaxtype(arg):
diff --git a/src/pybamm/solvers/jax_solver.py b/src/pybamm/solvers/jax_solver.py
@@ -6,6 +6,7 @@
 
 if pybamm.has_jax():
     import jax
+    import jax.extend
     import jax.numpy as jnp
     from jax.experimental.ode import odeint
 
@@ -62,9 +63,7 @@ def __init__(
         extra_options=None,
     ):
         if not pybamm.has_jax():
-            raise ModuleNotFoundError(
-                "Jax or jaxlib is not installed, please see https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-jaxsolver"
-            )
+            pybamm.raise_jax_not_found()
 
         # note: bdf solver itself calculates consistent initial conditions so can set
         # root_method to none, allow user to override this behavior
@@ -231,7 +230,7 @@ def _integrate(
         if model not in self._cached_solves:
             self._cached_solves[model] = self.create_solve(model, t_eval)
 
-        platform = jax.lib.xla_bridge.get_backend().platform.casefold()
+        platform = jax.extend.backend.get_backend().platform.casefold()
         if len(inputs) == 1:
             y = [self._cached_solves[model](y0_list[0], inputs[0])]
         elif platform.startswith("cpu"):
diff --git a/src/pybamm/util.py b/src/pybamm/util.py
@@ -362,6 +362,28 @@ def has_jax():
     )
 
 
+def is_macos_intel():
+    """Check if running on macOS with an Intel (x86_64) processor."""
+    import platform
+    import sys
+
+    return sys.platform == "darwin" and platform.machine() == "x86_64"
+
+
+def raise_jax_not_found():
+    """Raise an appropriate error when JAX is not available."""
+    if is_macos_intel():
+        raise ModuleNotFoundError(
+            "JAX is not available on macOS with Intel processors. "
+            "JAX dropped macOS x86_64 support in version 0.5.0. "
+            "JAX features require macOS with Apple Silicon (M-series), Linux, or Windows."
+        )
+    raise ModuleNotFoundError(
+        "Jax or jaxlib is not installed, please see "
+        "https://docs.pybamm.org/en/latest/source/user_guide/installation/gnu-linux-mac.html#optional-jaxsolver"
+    )
+
+
 def is_constant_and_can_evaluate(symbol):
     """
     Returns True if symbol is constant and evaluation does not raise any errors.
diff --git a/tests/unit/test_parameters/test_bpx.py b/tests/unit/test_parameters/test_bpx.py
@@ -477,6 +477,6 @@ def test_bruggeman_invalid_values_raise(self):
         bpx_obj["Parameterisation"]["Negative electrode"]["Porosity"] = 0  # Invalid
 
         with pytest.raises(
-            ValueError, match=r"math domain error"
-        ):  # Matches log(0) error
+            ValueError, match=r"math domain error|expected a positive input"
+        ):  # Matches log(0) error (message changed in Python 3.14)
             pybamm.ParameterValues.create_from_bpx_obj(bpx_obj)
diff --git a/tests/unit/test_solvers/test_base_solver.py b/tests/unit/test_solvers/test_base_solver.py
