[Bug]: Minimum JAX version is 0.4.16

[DavidMStraub]

PyBaMM Version

24.1

Python Version

3.10.12

Describe the bug

When a JAX version between 0.4.0 and 0.4.15 is installed and PyBaMM is installed without the [jax] optional (which can easily happen e.g. in an environment where pybamm[all]==24.1 is issued after pybamm[jax]==23.9 has been issued before), PyBaMM cannot be imported because pybamm.have_jax() only checks for a version greater than JAX_VERSION = "0.4" (so, returns True), but the import from jax.extend import linear_util as lu only starts working with JAX >= 0.4.16. (See #3671 where this was changed to silence a deprecation warning.)

I suggest to add a minimum JAX version (currently 0.4.16) and compare to that in have_jax.

Steps to Reproduce

python -m venv myenv
. myenv/bin/activate
pip install wheel
pip install jax==0.4.13 jaxlib==0.4.13 pybamm[all]==24.1
python -c 'import pybamm'

Relevant log output

Traceback (most recent call last):
    File "<string>", line 1, in <module>
    File "/home/straub/myenv/lib/python3.10/site-packages/pybamm/__init__.py", line 221, in <module>
      from .solvers.jax_bdf_solver import jax_bdf_integrate
    File "/home/straub/myenv/lib/python3.10/site-packages/pybamm/solvers/jax_bdf_solver.py", line 13, in <module>
      from jax.extend import linear_util as lu
ImportError: cannot import name 'linear_util' from 'jax.extend' (/home/straub/myenv/lib/python3.10/site-packages/jax/extend/__init__.py)

[arjxn-py]

Thanks @DavidMStraub for opening this, will look into this

[agriyakhetarpal]

We added a hard pin to the JAX version (to 0.4.27) in #4129. We usually don't unpin it or add bounds instead of pins because the JAX API is unstable and we find that there are breaking changes in every patch release. Please let us know if that helps, @DavidMStraub.

[DavidMStraub]

Thanks for the quick fix! Yes, I think this will work. However, I am wondering whether you should also have == instead of >= here in this case: https://github.com/pybamm-team/PyBaMM/blob/develop/pybamm/util.py#L357-L358

[agriyakhetarpal]

However, I am wondering whether you should also have == instead of >= here in this case: develop/pybamm/util.py#L357-L358

Ah, thanks for pointing this out! This function was used for the pybamm_install_jax entry point, which we replaced with the cleaner pip install pybamm[jax] starting with version 23.9, so it's not used nowadays. Based on the deprecation policy for three releases, we'll remove this command from the source code with version 24.9. I'll open a new issue about this for anyone wanting to take it after v24.5 is released.