Skip to content

Support Python 3.14, drop JAX on macOS Intel#5374

Merged
BradyPlanden merged 9 commits intomainfrom
increment-python-314
Feb 6, 2026
Merged

Support Python 3.14, drop JAX on macOS Intel#5374
BradyPlanden merged 9 commits intomainfrom
increment-python-314

Conversation

@BradyPlanden
Copy link
Member

@BradyPlanden BradyPlanden commented Feb 5, 2026

Description

Paired PR for pybamm-team/pybammsolvers#79

  • Bumps requires-python to <3.15, CasADi to 3.7.2, and JAX to >=0.7.0,<0.9.0 for Python 3.14 compatibility.
  • Adds 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.
  • Updates infra workflows to use 3.14, drops 3.10 from CI. Updates corresponding documentation

Fixes # (issue)

Type of change

Please add a line in the relevant section of CHANGELOG.md to document the change (include PR #)

Important checks:

Please confirm the following before marking the PR as ready for review:

  • No style issues: nox -s pre-commit
  • All tests pass: nox -s tests
  • The documentation builds: nox -s doctests
  • Code is commented for hard-to-understand areas
  • Tests added that prove fix is effective or that feature works

Bump requires-python to <3.15, CasADi to 3.7.2, and JAX to >=0.7.0,<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.
@BradyPlanden BradyPlanden requested a review from a team as a code owner February 5, 2026 14:31
BradyPlanden and others added 2 commits February 5, 2026 15:43
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.
@codecov
Copy link

codecov bot commented Feb 5, 2026

Codecov Report

❌ Patch coverage is 61.90476% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 98.30%. Comparing base (86531fb) to head (09a3108).
⚠️ Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
src/pybamm/util.py 25.00% 6 Missing ⚠️
src/pybamm/solvers/jax_bdf_solver.py 66.66% 1 Missing ⚠️
src/pybamm/solvers/jax_solver.py 66.66% 1 Missing ⚠️
Additional details and impacted files
@@           Coverage Diff           @@
##             main    #5374   +/-   ##
=======================================
  Coverage   98.29%   98.30%           
=======================================
  Files         328      328           
  Lines       28941    28952   +11     
=======================================
+ Hits        28449    28460   +11     
  Misses        492      492           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Add python_version >= '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).
Update all workflows, documentation, Dockerfile, ReadTheDocs config,
and pyproject.toml classifiers to include Python 3.14. Drop Python 3.10
from CI test matrices.
@BradyPlanden BradyPlanden merged commit 6d6f51a into main Feb 6, 2026
23 of 26 checks passed
@BradyPlanden BradyPlanden deleted the increment-python-314 branch February 6, 2026 16:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants