Unable to import jax

[gvijay1788]

Hello, I tried isntalling different version of jax and jaxlib and getting following error message while importing jax (import jax). Can anyone assist in this.

jax	Jaxlib	errror
0.3.13	0.3.5	AttributeError: module 'jax' has no attribute 'version'
jax-0.3.25		ImportError: cannot import name 'custom_ivjp' from 'jax.api' (/home/notebook/.virtualenvs/1267/lib/python3.7/site-packages/jax/api.py)
jax 0.1.75	jaxlib 0.3.5	AttributeError: 'str' object has no attribute 'name'

I am sharing screenshot of one of error. I have tried all different suggestion on google, unable to fix it. kindly suggest solution.

[jakevdp]

Is there a reason you prefer older JAX versions?

For JAX v0.3.0 or newer, the rule of thumb is to use the most recent jaxlib release that is not newer than the jax release. So JAX 0.3.15 goes with jaxlib 0.3.15, and jax 0.3.25 goes with jaxlib 0.3.25, and jax 0.3.19 goes with jaxlib 0.3.15 (because the next jaxlib release after 0.3.15 was 0.3.20). You can find more on this policy here: https://jax.readthedocs.io/en/latest/jep/9419-jax-versioning.html

If you need JAX releases older than 0.3.0, I would recommend first figuring out why that's required and trying to use a newer release instead. But if you still absolutely need an older jax version, I'd recommend finding the most recent jaxlib at the time of its release: the release history is listed at https://pypi.org/project/jax/#history and https://pypi.org/project/jaxlib/#history). For example, jax 0.2.10 was released on March 5 2021, at which point the most recent jaxlib release was 0.1.61 and so those versions should work together.

[gvijay1788]

Thanks jakevdp, I tried your suggstions but still getting err0r-

1. pip install jax==0.3.15 jaxlib==0.3.15
Error: import jax result following error-
AttributeError: module 'jax' has no attribute 'version'

2. pip install jax==0.3.25 jaxlib==0.3.25
Error: AttributeError: module 'jax' has no attribute 'version'

3. pip install jax==0.3.19 jaxlib==0.3.15
Error: AttributeError: module 'jax' has no attribute 'version'

[gvijay1788]

I am happy if newer version works well.

[jakevdp]

Thanks for the details. the error you're getting is strange, because jax.version is explicitly imported in the file above the line where this error is being raised (see https://github.com/google/jax/blob/jaxlib-v0.3.15/jax/_src/lib/__init__.py#L45). This makes me think that Python's import machinery is finding multiple copies of jax.

Do you by chance have multiple versions of jax in your Python import path? Or perhaps do you have a file named jax.py or a directory named jax in your path? I would start by digging in there, rather than trying to install different jax versions.

[jakevdp]

Actually, I bet this is the issue: if you're running in a jupyter or colab notebook, you need to restart your runtime after installing a new version of the package, otherwise the previously-imported version will be cached in sys.modules and you can encounter strange errors like this.

You can see this recommendation in the text in your screenshot, though it's easy to miss.

[gvijay1788]

Thanks for your suggestion. I am able to install with different version of python. Now M running following code and getting error while adding assertion.
Can you please help in this.

from jax.lib import xla_bridge
assert xla_bridge.get_backend().platform == 'gpu', 'Failed to locate GPU'

Can you please suggest.

[gvijay1788]

@jakevdp : Can you please help provide help on this ?

[jakevdp]

I would suggest not using pip installed CUDA for now. I recall hearing that there is some problem with that in the current release, but I don’t recall specifically what it is.

[gvijay1788]

and did we figure out solution or alternate for this problem? Could you please help me out to know that? May be get it from [hawkinsp] is you working together.

[jakevdp]

For help with installing CUDA, I'd recommend following the CUDA docs; e.g. https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html

[jakevdp]

FYI, there's some info here on why pip-installed CUDA does not work with JAX: #13008 (comment) We're hoping NVIDIA is able to fix that soon so we can streamline the jaxlib-gpu install process.