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

Question

对于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一起使用，然后按照注释的建议，通过克隆它们的存储库并按照此处给出的说明从源代码构建它，但显示了相同的错误消息。

Answer 1

我也遇到了类似的问题。由于我已经安装了Anaconda，不想用Anaconda+miniconda+自制版本的python和包管理之类的东西来扰乱我的空间，所以我四处寻找一个简单的解决方案。最终对我起作用的是首先卸载jax和jaxlib，然后通过conda-forge直接安装jax和jaxlib：

pip uninstall jax jaxlib
conda install -c conda-forge jaxlib
conda install -c conda-forge jax

Answer 2

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

感谢Noah在第5501期中提到，您可以使用jax和jaxlib的早期版本，就我的目的而言，jaxlib==0.1.60和jax==0.2.10工作得很好！

Answer 3

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

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

Answer 4

截至目前(2022年1月(，jax可用于M1 Mac。确保卸载jax和jaxlib，然后通过pip:安装新软件包

pip install --upgrade jax jaxlib

之后，您可以毫无问题地使用jax。

--编辑--我在一台具有以下规格的机器上运行：

ProductName:    macOS
ProductVersion: 12.1
BuildVersion:   21C52

并且与CCD_ 13在CCD_。