python - 如何处理 ImportError : cannot import name 'pytree' from 'jaxlib' of using jax?

Question

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

Successfully installed jax-0.1.68 jaxlib-0.1.67+cuda11

但是在我运行我的项目之后，它显示了 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)

这个问题可能来自运行项目的不兼容的 jax/jaxlib 版本吗？
如果不是，如何处理？

Answer 1

看来您正在导入比您在问题中报告的旧得多的 jax 版本； jax.lib未尝试导入 pytree来自 jaxlib从 0.2.8 版开始。
这可能表明您正在运行 pip install在与您用来执行代码的环境不同的环境中。
假设您在命令提示符下工作，请尝试以下操作:

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

(您可以将两行中的 python 替换为您正在使用的任何 python 可执行文件)
如果您在 Jupyter 中使用不同的内核，这个答案可能会帮助您了解如何继续: Running Jupyter with multiple Python and IPython paths