jax Import Problem - partially initialized module 'jax' has no attribute 'version' (most likely due to a circular import)**

[LukeAndreesen]

Hello! I'm attempting to use Jax for a project, but am having trouble importing it. I installed Jax using pip install -U "jax[cpu]" in the directory of my project. I then attempted to import jax in a jupyter notebook, and received the following error:

AttributeError: partially initialized module 'jax' has no attribute 'version' (most likely due to a circular import)
---> 76 jax_version=jax.version.version,

I've tried to uninstall and reinstall, as well as install in a virtual environment, but consistently am running into this issue when i run:
import jax
in a notebook. Any troubleshooting help would be greatly appreciated! Thanks so much.

[dfm]

I have seen this happen because of mismatched jax and jaxlib versions in the past. Could you share which platform you're on (mac, linux, etc.?) and the output of pip freeze?

[dfm]

Strange! Well my best guess is that somehow the jupyter environment is finding conflicting versions somehow. You might want to see if it works in the Python REPL, or start digging into PYTHONPATH environment variables, but I'm not sure I can make too many concrete suggestions since I can't reproduce the issue on my mac. Please do report back if you find the culprit!

[LukeAndreesen]

Will do! Thank you so much for your help, I'll follow up if I get to the source of the problem!

[jakevdp]

Is it possible that your jupyter notebook had an existing runtime that imported an older jax before you updated the jax version on your system? That's another way to get the kind of version skew that would lead to this error. You could fix that by restarting the Jupyter kernel and then running your code again.

[jakevdp]

I answered a similar question here with some more ideas.

[LukeAndreesen]

Thank you all for your help! I eventually found some packages in the form /MathJax/jax/output/HTML-CSS/jax.js that were somewhere in my computer - removing these did the trick!