[Bug]: Proposal: Check JAX Version Compatibility Before Calling JAX Functionalities

[thegialeo]

PyBaMM Version

25.12.2

Python Version

3.12

Describe the bug

First of all, thank you for the great work on PyBaMM. I’d like to propose an improvement regarding how JAX availability is handled.

Summary

Currently, has_jax() in util.py only checks whether jax is installed, but not which version is installed. This can lead to unexpected import-time failures when an incompatible JAX version is present in the environment.

If PyBaMM is installed with the pybamm[jax] extra, dependency resolvers (e.g., Poetry, etc.) correctly enforce the required JAX version constraints. However, when installing PyBaMM without the JAX extra, there is an assumption (both by users and dependency resolvers) that PyBaMM is completely independent of JAX.

In practice, this is not entirely true. If JAX is already installed in the environment (e.g., via a transitive dependency from another library or as a direct dependency), PyBaMM may still attempt to call JAX functionality and break the existing codebase.

Minimal Reproduction

conda create -n my_env python=3.12
conda activate my_env
pip install pybamm
pip install jax

Then:

import pybamm

This results in:

AttributeError: module 'jax.lib' has no attribute 'xla_bridge'

This appears to happen because PyBaMM assumes an older JAX API (jax>=0.4.36,<0.7.0), while newer JAX releases introduce changes that may not be fully compatible with that expected API.

Problem Description

• Users installing PyBaMM without [jax] reasonably expect no interaction with JAX.
• If JAX is present in the environment (even indirectly), PyBaMM may attempt to use it without verifying compatibility.
• The resulting error is non-obvious and can take time to diagnose, especially for new users unfamiliar with PyBaMM’s optional JAX integration.
• This creates unnecessary friction and can give the impression that PyBaMM is unstable, when the root cause is simply a version mismatch.

In my case, switching to the main branch resolved the issue because the JAX version requirement has been bumped there. However, I think it would still be beneficial to explicitly check the JAX version and gracefully ignore JAX support if the installed version is incompatible.

Proposal

Instead of only checking whether JAX is installed, has_jax() in util.py could additionally validate that the installed version satisfies PyBaMM’s supported version range.

If not compatible, PyBaMM could either:

• Disable JAX support (i.e. treat it as has_jax() == False)
• Raise a clear and informative error message explaining the version mismatch, explicitly stating that PyBaMM has a JAX version dependency whenever JAX is present in the environment even if PyBaMM was installed without the [jax] extra and clearly informing the user that a JAX installation was detected and a specific version of JAX is assumed.

As a user, I would prefer the former option.

This would:

• Make the optional dependency behave more like a truly optional dependency
• Avoid hard import-time crashes
• Provide clearer guidance to users

Steps to Reproduce

conda create -n my_env python=3.12
conda activate my_env
pip install pybamm
pip install jax

Then:

import pybamm

Relevant log output

AttributeError: module 'jax.lib' has no attribute 'xla_bridge'

[agriyakhetarpal]

Thanks, @thegialeo. We previously had functionality a few years ago that did exactly this – i.e., check for the JAX version and raise an informative error if it did not match, but we dropped it for some reason. We should indeed bring this back. I believe we had loosened the requirement as JAX became more stable by the day, but the error you encountered suggests we should still be cautious. This should be a good issue to take up for someone exploring the codebase to learn, but by all means, please feel free to put together a PR yourself if you wish to contribute!

[thegialeo]

I’m not super familiar with the PyBaMM codebase yet, but the implementation itself should not be too complicated. The bigger question to me is which approach you’d prefer long term. Since you’ll ultimately be maintaining this, I think it makes sense for you to decide on the direction. Once that’s clear, I’m happy to put together a PR and iterate from there.

I can think of three possible approaches:

1. Hard-code a version check
   We could read the installed jax/jaxlib versions via importlib.metadata.version() and compare them against a supported range. This is simple and doesn’t add any new dependencies. The downside is that the version range would be duplicated (once in pyproject.toml, once in code), which creates a risk of drift if someone updates the dependency constraint but forgets to update the check. That would probably require an extra test to keep things in sync.

2. Read the version requirement from pyproject.toml
   Instead of duplicating the range, we could try to read it from project metadata and validate against that. Cleaner in principle (single source of truth), but more complex to implement. It likely means adding some packaging/parsing logic just to support this check, which might be overkill.

3. Don’t check versions explicitly, just catch failures
   Another option would be to stop proactively validating versions and instead just guard JAX usage more defensively: try importing JAX if anything goes wrong (import error, attribute error, etc.), emit a warning and treat JAX as unavailable i.e. return False. This is the simplest approach and avoids version duplication entirely. It also protects against breakages that aren’t strictly version-related. It also has the fewest moving parts, meaning fewer potential points of failure and less extra code that needs to be maintained over time. The trade-off is that incompatibilities are detected reactively rather than upfront.

The more I think about it, the more I feel like option 3 might actually be sufficient. Since JAX support is optional, maybe the most important thing is that PyBaMM never hard-crashes because of it. If JAX is present but unusable for any reason, treating it as “not available” and warning the user might be the most robust behavior.

But I’m very open to whichever approach you think fits best with the project’s direction. Once you’ve decided, I’ll go ahead and draft a PR.

[agriyakhetarpal]

Thanks! I would suggest we choose Option 1, which is what we had previously. Regarding the risk of drift, we previously noted it in a comment near the JAX dependency in pyproject.toml, saying "When updating the versions here, also update the checks in file X", and similarly in the other file to suggest updating pyproject.toml. I would have been fine with Option 3, too, but the issue is that most people tend to ignore warnings, so Option 1 provides fewer surprises that way.

[thegialeo]

I’ll proceed with option 1, along with cross-referenced comments in both pyproject.toml and the relevant source file to help prevent drift between them.

Just to make sure we’re aligned: regardless of whether we choose option 1, 2, or 3, PyBaMM should never raise a hard error due to a JAX version mismatch. If the installed jax/jaxlib versions are unsupported, we should treat JAX as unavailable, emit a warning, and have has_jax() return False.

For users who explicitly require JAX support in PyBaMM, installing via pybamm[jax] should already ensure compatible versions through the dependency resolver. In that case, managing JAX version constraints is not PyBaMM’s responsibility as a library.

The primary scenario this change addresses is users who do not intend to use JAX with PyBaMM, install PyBaMM without the jax extra, but happen to have JAX in their environment (e.g., for other numeric computations or as a transitive dependency). In those cases, PyBaMM should not fail on import or raise an exception. Instead, it should behave as though JAX is unavailable, ideally with a clear warning explaining why.

If this matches your expectations, I’ll put together a PR later this week.