如何处理ImportError:不能从jax的'jaxlib‘导入名称'pytree’？

Question

我根据https://github.com/google/jax#installation使用pip安装了jax和jaxlib

Successfully installed jax-0.1.68 jaxlib-0.1.67+cuda11

但是在我运行了我的项目之后，它显示了我的恐惧：

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)

这个问题是否来自运行项目时不兼容的jax/jaxlib版本？如果没有，如何处理呢？

Answer 1

似乎您正在导入一个比您在问题中报告的要早得多的jax版本；自0.2.8版本以来，jax.lib一直没有尝试从jaxlib导入pytree。

这可能表明您运行的pip install环境与用于执行代码的环境不同。

假设您在命令提示符下工作，请尝试以下操作：

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

(在这两行中，您都可以用您正在使用的python可执行文件替换python )

如果您在木星中使用不同的内核，这个答案可能会帮助您理解如何进行：运行具有多个Python和IPython路径的木星