AttributeError: partially initialized module 'jax' has no attribute 'version' (most likely due to a circular import)

[lexie-demarco]

I am trying to use jax.numpy and am getting the error "AttributeError: partially initialized module 'jax' has no attribute 'version' (most likely due to a circular import)" Nothing I have tried worked and would appreciate any help.

[jakevdp]

This typically means you have some kind of mismatch between the version of jax and the version of jaxlib in your environment. You can check those by running:

pip list | grep jax

To fix it, try updating to the most recent jax/jaxlib version:

pip install -U jax jaxlib

(If you need GPU support, see the install guide)

Also be certain to restart your Colab runtime any time you update packges, otherwise cached package versions may be used and could cause the error you're seeing.

[GeorgiadeG]

Did you restart your notebook session after installing new versions in the runtime?

I had the same error, update both jax and jaxlib exactly the same as shown by OP, restart my notebook and the error doesn't change. Any ideas of what else might be wrong?

[jakevdp]

The error here indicates that the jax and jaxlib versions are out-of-sync. This can be caused by having mismatched versions installed, or by updating versions without restarting your jupyter or colab runtime.

Are you certain you're importing the jax/jaxlib versions you think you are? Especially when working in Jupyter notebooks, it's possible that you're installing the packages in a different environment than your notebook runtime. There's a decent writeup of this issue here: https://stackoverflow.com/questions/39007571/running-jupyter-with-multiple-python-and-ipython-paths/39022003#39022003

[dmilakovic]

Hi.
I have the same error.

It seems to me that the issue stems from the fact that jax.version is not defined before provided to function check_jaxlib_version in line 74 of the file jax/_src/lib/__init__.py, As a result, the function check_jaxlib_version fails to run at all.

Reading previous responses, I confirmed that the jax and jaxlib versions match so that is unlikely to be the culprit. Yet, the problem persists. Other information that may be useful:
jax 0.4.10
jaxlib 0.4.10
jaxopt 0.6

Python 3.9.12
MacOS 12.2.1

I am not using jupyter or other collaborative tools.

Many thanks!

[jakevdp]

The only time I've seen this error is when JAX and jaxlib are out of sync. If you believe you have a different issue, I'd suggest opening a new issue with more details. Long comment threads on old discussions tend to fall under the radar.

[LiamLau23]

but why i run the demand "pip list | grep jax" is wrong?

[Gaurav17Joshi]

I was getting the same issue, and I uninstalled and used conda forge to reinstall jax and jaxlib conda install -c conda-forge jax, conda install -c conda-forge jaxlib. My versions for both were 0.4.12.

[ingomueller-net]

I also ran into this issue after installing the wheel I built for jaxlib and then installing jax with pip install -e .. When I installed jax with pip install . (i.e., not as editable), then it the error went away.

[hepnerb]

Switching from Anaconda on Windows to WSL2 solved the issues for me. Not a full fix but a quick one.

[rajkk1]

I had the same problem, but then I uninstalled jax and jaxlib, and re-installed GPU version of jax with pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html (and didn't reinstall jaxlib) and that fixed the problem somehow 🤷🏾

[sinodine]

I tried that, but it didn't work for me.
What worked was:

• unistall jax and jaxlib
• pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
• pip install -U "jax[cpu]"
  The last command installed the jax 0.4.6 installed with the second line and reinstalled on it's own jax-0.4.28 and jaxlib-0.4.28.