Trying to import jax raises ImportError

[raghuramshankar]

Trying to import Jax raises ImportError:

import jax

This gives me the error:

File "/home/user/.local/lib/python3.9/site-packages/jax/_src/api.py", line 52, in <module>
    from ..tree_util import (tree_map, tree_flatten, tree_unflatten, tree_structure,
ImportError: cannot import name 'Partial' from 'jax.tree_util' (/home/user/.local/lib/python3.9/site-packages/jax/tree_util.py)

Python and jax version:

python --version
Python 3.9.6

pip show jax 
Name: jax
Version: 0.2.19

My system info:
Linux Manjaro (Arch Linux) with the latest kernel.

uname -r
5.13.12-1-MANJARO

Adding Partial to the list in /home/user/.local/lib/python3.9/site-packages/jax/tree_util.py line 40 solved the issue for me and I am able to successfully import jax.
I also noticed that partial is available in the list but replacing partial with Partial seems to raise more errors for me.

[jakevdp]

Hi - thanks for the report!

I'm a bit confused about what's going on, because you mention that you're using JAX 0.2.19, and you mention adding Partial at line 40 of jax/tree_util.py, but if I look at that line within the 0.2.19 source tree Partial is already there: https://github.com/google/jax/blob/jax-v0.2.19/jax/tree_util.py#L40

I would do two things: first, check the output of this:

$ python -m pip show jax

calling pip this way assures you're using the pip associated with your python executable, which is not guaranteed if you just call pip alone.

If that also shows jax v 0.2.19, then I assume your installation has been corrupted somehow. I'd then run

$ python -m pip uninstall jax jaxlib
$ python -m pip install jax[cpu]

(or the equivalent GPU installation command if relevant)

[ioannis12]

hello,
I think I have the same problem, when I try to import jax I get an error message "kernel died, restarting kernel".
I use spyder 4.1.0.

I tried uninstalling and installing different versions, but it doesn't seem to work

[raghuramshankar]

Hi,

Thank you for your reply.
I am using Jax as a solver for Pybamm simulations. I just found that the issue was with the installation of the Pybamm library, and not with the Jax library. Here is the issue that was raised over there regarding the same.
Sorry for not catching this earlier, I did not expect it to be this way.
Thanks!

Regards
Raghuram

[jakevdp]

Thanks! I'll close the issue then. Feel free to open another if you run into any other problems!

[RylanSchaeffer]

I received the same error from jax 0.2.21:

  File "/om2/user/rylansch/FieteLab-RIBP/ribp/lib/python3.7/site-packages/numpyro/infer/hmc.py", line 8, in <module>
    from jax import device_put, lax, partial, random, vmap
ImportError: cannot import name 'partial' from 'jax' (/om2/user/rylansch/FieteLab-RIBP/ribp/lib/python3.7/site-packages/jax/__init__.py)
python-BaseException

[RylanSchaeffer]

I think the problem is with numpyro (pyro-ppl/numpyro#1165)

[jakevdp]

jax.partial was removed in version 0.2.21 (see the Changelog). You should import functools.partial instead, or install version 0.2.20 if you depend on other packages that attempt to import partial from the jax namespace.

[yukiweng62]

I received a similar error from jax 0.4.11, could u help me:

File C:\ProgramData\anaconda3\Lib\site-packages\flax\linen\linear.py:30
from jax import ShapedArray

ImportError: cannot import name 'ShapedArray' from 'jax' (C:\ProgramData\anaconda3\Lib\site-packages\jax_init_.py)

[jakevdp]

jax.ShapedArray was removed in JAX v0.4.14 after being deprecated for several releases (see the JAX v0.4.14 changelog).

I would suggest updating flax to a more recent version that is compatible with recent JAX versions.

[yukiweng62]

Thanks a lot!!!

[aiyb1314]

ImportError: cannot import name 'abstract_arrays' from 'jax' (/home/xuj/anaconda3/envs/removal/lib/python3.10/site-packages/jax/init.py)
How to solve??

[jakevdp]

jax.abstract_arrays was deprecated in jax v0.4.12 and removed in JAX v 0.4.17.

To fis the error you're seeing, you can avoid attempting to import jax.abstract_arrays. If you're using another package that imports this deprecated module, you could try updating that other package, or if that's not possible, install JAX v0.4.16 or older.