使用 m1 芯片在 mac 上导入 jax 失败

问题描述

对于 python 3.8.8 并使用新的 mac air（带有 m1 芯片），在 jupyter 笔记本和 python 终端中，import jax 引发此错误

Python 3.8.8 (default,Apr 13 2021,12:59:45)
[Clang 10.0.0 ] :: Anaconda,Inc. on darwin
Type "help","copyright","credits" or "license" for more information.
>>> import jax
Traceback (most recent call last):
  File "<stdin>",line 1,in <module>
  File "/Users/steve/Documents/code/jax/jax/__init__.py",line 37,in <module>
    from . import config as _config_module
  File "/Users/steve/Documents/code/jax/jax/config.py",line 18,in <module>
    from jax._src.config import config
  File "/Users/steve/Documents/code/jax/jax/_src/config.py",line 26,in <module>
    from jax import lib
  File "/Users/steve/Documents/code/jax/jax/lib/__init__.py",line 63,in <module>
    cpu_feature_guard.check_cpu_features()
RuntimeError: This version of jaxlib was built using AVX instructions,which your cpu and/or operating system do not support. You may be able work around this issue by building jaxlib from source.

我怀疑是 m1 芯片造成的。

我尝试将 jax 与 pip install jax 一起使用，然后按照评论的建议从源代码构建它，方法是克隆他们的存储库并按照 here 给出的说明进行操作，但显示相同的错误消息。>

解决方法

JAX 尚未提供与 M1 芯片兼容的预构建 jaxlib 轮。我所知道的关于在 M1 上构建 jaxlib 的最佳信息来源可能是这个 github 问题：https://github.com/google/jax/issues/5501，它也跟踪改进这种支持。

希望 M1 支持将在不久的将来得到改善，但堆栈上下的科学计算基础设施需要一段时间才能赶上新芯片的要求。

,

谢谢 @jakevdp 我查看了您链接的问题并找到了解决方法：

感谢 Noah 在 issue #5501 中提到您可以只使用先前版本的 jax 和 jaxlib，就我而言，jaxlib==0.1.60 和 jax==0.2.10 可以正常工作！