How to deal with ImportError: cannot import name 'pytree' from 'jaxlib' of using jax?

user_17643340 0
Viewed 1053

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?

jax importerror ubuntu python
1 Answers

It appears that you are importing a much older jax version than you report in the question; jax.lib has not attempted to import pytree from jaxlib since version 0.2.8.

This probably indicates that you are running pip install in a different environment than the one you're using to execute code.

Assuming you are working at the command prompt, try this instead:

$ python -m pip install jax jaxlib
$ python -c "import jax; print(jax.__version__)"

(where you can replace python in both lines with whatever python executable you are using)

If you're working in Jupyter with different kernels, this answer might help you understand how to proceed: Running Jupyter with multiple Python and IPython paths

user_29378310
Related