What is the compatible version between CUDA, pybamm and jax

[lcwxz1989]

PyBaMM Version

24.9

Jax Version

0.4.27

Python Version

3.10

CUDA Version

12.6

Describe the bug

The question is Jax 0.4.27 version has bug, the Info is below, and Jax fix the bug in version 0.4.28. But Pybamm 24.9 can only support 0.4.27, so I need Pybamm can support Jax version newer than 0.4.27

Info

'+ptx86' is not a recognized feature for this target (ignoring feature)

Image

Steps to Reproduce

pip install -U "jax[cuda12]==0.4.27"
pip install -U pybamm==24.9.0
sudo apt-get install -y nvidia-open

Script

this nootbook demo

https://docs.pybamm.org/en/latest/source/examples/notebooks/solvers/idaklu-jax-interface.html

Relevant log output

'+ptx86' is not a recognized feature for this target (ignoring feature)

No response

[kratman]

@jsbrittain I am guessing this is a question for you

[jsbrittain]

@lcwxz1989 PyBaMM comes with the default Jax (CPU) install by default. You can upgrade this to support cuda on your system (cuda-12 by the look of it in this case) by:

pip install --upgrade "jax[cuda12]"

If you now re-launch Python and import pybamm and jax, the GPU warning will no longer be displayed and jax.devices() should show cuda as an available device. Details on installing Jax with support for various device can be found here: https://jax.readthedocs.io/en/latest/installation.html.

[lcwxz1989]

the env is so confused, pybamm 24.9.0 only compare with jax 0.4.27. but jax 0.4.27 will encountered error pic1.
jax fix the "'+ptx86' is not a recognized feature for this target (ignoring feature)" in 0.4.28. But if I update jax to 0.4.28, I got error pic2 which pybamm can not find jax. So is there a pybamm with jax for gpu version?

pic1:

pic2:

[agriyakhetarpal]

I'll convert this to a discussion since it doesn't seem like there is a bug on our end – more details can be shared there. A CUDA-enabled JAX distribution is required for GPU support.

[lcwxz1989]

BTW, I can understand this project is focus on exprimental. Therefore, strictly speaking, efficiency is not the goal of the project. But this is really a import feature when we do exprimental.

[kratman]

@lcwxz1989 Part of the difficulty is that many features are provided by volunteer developers. This makes it take a bit longer for features like GPU support to become fully stable. There are some researchers working on improving the solver and GPU support, so hopefully it will continue to improve.

If efficiency is a major concern, I would suggest checking out the IDAKLUSolver. It runs on CPUs, but is significantly faster than the CasadiSolver which is the default in PyBaMM.

[lcwxz1989]

@lcwxz1989 Part of the difficulty is that many features are provided by volunteer developers. This makes it take a bit longer for features like GPU support to become fully stable. There are some researchers working on improving the solver and GPU support, so hopefully it will continue to improve.

If efficiency is a major concern, I would suggest checking out the IDAKLUSolver. It runs on CPUs, but is significantly faster than the CasadiSolver which is the default in PyBaMM.

Thanks a lot, I will try.
And for this discussion, would you try to up jax 0.4.28? I can help to make a test

[lcwxz1989]

Hi @lcwxz1989, I think this is because we don't support 0.4.28 yet. We hardcode the JAX version here:

PyBaMM/src/pybamm/util.py

Lines 14 to 17 in f20a00e

	# Versions of jax and jaxlib compatible with PyBaMM. Note: these are also defined in
	# the extras dependencies in pyproject.toml, and therefore must be kept in sync.
	JAX_VERSION = "0.4.27"
	JAXLIB_VERSION = "0.4.27"

, and any other version will make pybamm.has_jax() return False as an output.
this is done in order to prevent breaking changes or unexpected bugs from creeping in with our implementation of the JaxSolver. Would you like to propose bumping the version? Based on https://github.com/jax-ml/jax/blob/main/CHANGELOG.md#jaxlib-0428-may-9-2024, I think you'll need to use a lower CUDA version for now. We don't have a JAX-for-GPU version so far, since we don't (and can't) regularly test it because of the absence of GPU runners in our CI infrastructure.

If so, I can build by editing the version to 0.4.28 , and pybamm.has_jax() will return ture, right?

[lcwxz1989]

Hi @lcwxz1989, I think this is because we don't support 0.4.28 yet. We hardcode the JAX version here:

PyBaMM/src/pybamm/util.py

Lines 14 to 17 in f20a00e

	# Versions of jax and jaxlib compatible with PyBaMM. Note: these are also defined in
	# the extras dependencies in pyproject.toml, and therefore must be kept in sync.
	JAX_VERSION = "0.4.27"
	JAXLIB_VERSION = "0.4.27"

, and any other version will make pybamm.has_jax() return False as an output.
this is done in order to prevent breaking changes or unexpected bugs from creeping in with our implementation of the JaxSolver. Would you like to propose bumping the version? Based on https://github.com/jax-ml/jax/blob/main/CHANGELOG.md#jaxlib-0428-may-9-2024, I think you'll need to use a lower CUDA version for now. We don't have a JAX-for-GPU version so far, since we don't (and can't) regularly test it because of the absence of GPU runners in our CI infrastructure.

If we have hard code "0.4.27", I think we have a compaitble cuda verison? Can you share about cuda version?

[lcwxz1989]

If we have hard code "0.4.27", I think we have a compaitble cuda verison? Who can share about cuda version?