在 Mac 上使用 M1 芯片时导入 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芯片引起的。

我尝试使用pip install jax来使用jax，然后按照评论中建议的方式从源代码构建它，通过克隆他们的存储库并遵循这里给出的说明，但是仍然显示相同的错误信息。

- dcxst

Answer 1

我遇到了类似的问题。由于我已经安装了Anaconda，不想在我的空间中弄乱Anaconda+miniconda+homebrew版本的Python和包管理等内容，因此我寻找了一个简单的解决方案。最终对我有用的是先卸载jax和jaxlib，然后直接通过conda-forge安装jax和jaxlib：

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

- Nicholas G Reich

---

不是问题的重点，但您能否使用trax或者只想使用jax？ - arturomp

我刚刚在看jax，还没有看过trax。 - Nicholas G Reich

@NicholasGReich，我刚好在寻找这个确切的问题的答案，世界真小啊。 - Kostis

有没有非conda解决方案？我希望能够安装到通过virtualenv（https://virtualenv.pypa.io/en/latest/）创建的虚拟环境。 - undefined

Answer 2

感谢@jakevdp，我查看了您提供的问题并找到了解决方法:

感谢Noah在问题#5501中提到，您可以使用先前版本的jax和jaxlib，对于我的目的，jaxlib==0.1.60和jax==0.2.10完全可以使用！

- dcxst

---

jax现在可用于M1 Mac，我运行没有问题。 - emil

没错，@emil。这里有确认信息：https://github.com/google/jax/issues/5501#issuecomment-955590288 - 下一个版本发布时间不确定，但至少现在可以正式安装了。 - arturomp

Answer 3

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

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

- jakevdp

Answer 4

截至目前（2022年1月），jax已经支持M1 Macs。请确保卸载jax和jaxlib，然后通过pip安装新的软件包：

pip install --upgrade jax jaxlib

之后，您就可以无障碍地使用jax了。

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

ProductName:    macOS
ProductVersion: 12.1
BuildVersion:   21C52

使用 conda 环境中的 Python 3.9.6。

- emil

---

你使用的是哪个MacOS和Python版本？ - Blade

@Blade，我已经更新了答案以增加清晰度。 - emil

卸载了jax和jaxlib，再用pip install --upgrade jax jaxlib安装后，我在使用Trax时仍然遇到了这个错误:RuntimeError: This version of jaxlib was built using AVX instructions, which your CPU and/or operating system does not support. You may be able work around this issue by building jaxlib from source. 这是相同的操作系统版本和构建，Python 3.9.7在一个conda环境中。但在某个特定点后仍然出现了这个错误。 - arturomp

@arturomp请看我刚刚发布的答案。我通过直接从conda-forge安装来解决了这个问题。 - Nicholas G Reich

谢谢！我认为我做了类似的事情：https://dev59.com/y8Pra4cB1Zd3GeqPgG9m#70815865/ - arturomp