ImportError: cannot import name 'pytree' from 'jaxlib'

[lzhhh93]

I have installed jax and jaxlib using pip according to https://github.com/google/jax#installation:

Successfully installed jax-0.1.68 jaxlib-0.1.67+cuda11

But after i ran my project, it showed Importerror:

Traceback (most recent call last):
  File "...", line 1, in <module>
    from jax import jit, jacfwd, jacrev, hessian, lax
  File "...", line 16, in <module>
    from .api import (
  File "...", line 38, in <module>
    from . import core
  File "...", line 30, in <module>
    from . import dtypes
  File "...", line 31, in <module>
    from .lib import xla_client
  File "...", line 51, in <module>
    from jaxlib import pytree
ImportError: cannot import name 'pytree' from 'jaxlib' (/home/control/.local/lib/python3.7/site-packages/jaxlib/__init__.py)

Could this problem comes from uncompatible jax/jaxlib version for running the project?
If not, how to deal with it?

Many thanks.

[hawkinsp]

It sounds to me like something isn't working with your jaxlib installation. pytree is part of jaxlib. Perhaps try uninstalling and reinstalling jaxlib? What Python version are you using?

[lzhhh93]

Thanks for your answer. I have uninstalled both jax and jaxlib from pycharm and installed again, but got the same result.
i am using python3.7:

$ pip install --upgrade jax==0.1.68 jaxlib==0.1.67+cuda111 -f https://storage.googleapis.com/jax-releases/jax_releases.html
Defaulting to user installation because normal site-packages is not writeable
Looking in links: https://storage.googleapis.com/jax-releases/jax_releases.html
Collecting jax==0.1.68
  Using cached jax-0.1.68-py3-none-any.whl
Collecting jaxlib==0.1.67+cuda111
  Using cached https://storage.googleapis.com/jax-releases/cuda111/jaxlib-0.1.67%2Bcuda111-cp37-none-manylinux2010_x86_64.whl (194.6 MB)
Requirement already satisfied: numpy>=1.12 in /usr/local/lib/python3.7/dist-packages (from jax==0.1.68) (1.20.0)
Requirement already satisfied: absl-py in ./.local/lib/python3.7/site-packages (from jax==0.1.68) (1.0.0)
Requirement already satisfied: opt-einsum in ./.local/lib/python3.7/site-packages (from jax==0.1.68) (3.3.0)
Requirement already satisfied: scipy in ./.local/lib/python3.7/site-packages (from jaxlib==0.1.67+cuda111) (1.7.3)
Requirement already satisfied: flatbuffers<3.0,>=1.12 in ./.local/lib/python3.7/site-packages (from jaxlib==0.1.67+cuda111) (1.12)
Requirement already satisfied: six in /usr/lib/python3/dist-packages (from absl-py->jax==0.1.68) (1.11.0)
Installing collected packages: jaxlib, jax
Successfully installed jax-0.1.68 jaxlib-0.1.67+cuda111

Edit: sori, the ImportError occurs when i install jax-0.1.68 jaxlib-0.1.67+cuda111, which is used to developed the project at the beginning, not jax-0.2.26, i have posted it wrong.

[jakevdp]

I noticed I answered the same question here: https://stackoverflow.com/questions/70331558/how-to-deal-with-importerror-cannot-import-name-pytree-from-jaxlib-of-using/70336158?noredirect=1#comment124335932_70336158

From your comments there, it sounds like you want to install a very old version of jax (0.1.68), and you're finding that a much newer version of jaxlib (0.1.67) is incompatible. This is to be expected.

If you need to use jax 0.1.68, you'll probably need jaxlib somewhere around 0.1.50, which you can install this way:

$ pip install jax==0.1.68 jaxlib==0.1.50 -f https://storage.googleapis.com/jax-releases/jax_releases.html

Note that the jax-releases address is necessary because this jaxlib version is no longer on PyPI due to its storage constraints.

If you need a GPU-compatible jaxlib, you may be able to use jaxlib-0.1.52, which is the first version for which we published GPU wheels (I haven't checked its compatibility with jax 0.1.68). That said, you may have system compatibility issues with those wheels because they were only built for CUDA 10.0 and 10.1 (Again, the versions you're asking about are very old).

Overall, I would recommend using newer jax/jaxlib versions if possible.

[lzhhh93]

sorry about cross posting. jaxlib 0.1.50 works very well and the code is running. Thank you very much for your help!