importing jax fails on mac with m1 chip

Question

For python 3.8.8 and using the new mac air (with the m1 chip), in jupyter notebooks and in python terminal, import jax raises this error

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.

I suspect it occurs because of the m1 chip.

I tried using jax with pip install jax, then I built it from source as suggested by the comment, by cloning their repository and following the instructions given here, but the same error message shows.

score 5 · Answer 1 · answered Jul 12 '21 at 10:18

5

Thanks @jakevdp I looked at the issue you linked and found a workaround :

Thanks to Noah who mentioned in issue #5501 that you could just use a previous version of jax and jaxlib, for my purposes jaxlib==0.1.60 and jax==0.2.10 work just fine!

answered Jul 12 '21 at 10:18

dcxst

173
3
8

2

`jax` is now available for M1 macs; runs fine for me – emil Jan 12 '22 at 08:08
1

True, @emil. Confirmation here: https://github.com/google/jax/issues/5501#issuecomment-955590288 - unclear when next release is, but at least it can be installed now somewhat officially – arturomp Jan 20 '22 at 22:10

score 5 · Answer 2 · answered Feb 08 '22 at 17:50

5

I had a similar problem. Since I already had Anaconda installed and didn't want to clutter up my space with Anaconda + miniconda + homebrew versions of python and package management and whatever, I hunted around for a simple solution. What ended up working for me was first uninstalling jax and jaxlib and then installing jax and jaxlib via conda-forge directly:

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

answered Feb 08 '22 at 17:50

Nicholas G Reich

1,028
10
21

not the point of the question, but are you able to use with `trax` or were you looking to use just `jax`? – arturomp Feb 08 '22 at 17:53
1

I was just looking at `jax`. hadn't looked at `trax`. – Nicholas G Reich Feb 08 '22 at 20:20
1

@NicholasGReich, was just looking for an answer to this exact question, small world. – Kostis Apr 14 '23 at 18:00

score 2 · Answer 3 · answered Jul 10 '21 at 14:50

JAX does not yet provide pre-built jaxlib wheels that are compatible with M1 chips. The best source of information I know on building jaxlib on M1 is probably this github issue: https://github.com/google/jax/issues/5501, which also tracks improving this support.

Hopefully M1 support will be improved in the near future, but it's taking a while for the scientific computing infrastructure up and down the stack to catch up with the requirements of the new chips.

emil · Answer 4 · 2022-01-15T10:52:56.797

0

As of now (January 2022), jax is available for M1 Macs. Make sure to uninstall jax and jaxlib and then install the new packages via pip:

pip install --upgrade jax jaxlib

Afterwards, you can use jax without problems.

--Edit-- I am running on a machine with the following specs:

ProductName:    macOS
ProductVersion: 12.1
BuildVersion:   21C52

and with Python 3.9.6 within a conda environment.

edited Jan 15 '22 at 10:52

answered Jan 12 '22 at 08:07

emil

194
1
11

What MacOS and python versions are you using? – Blade Jan 12 '22 at 20:04
@Blade, I have updated the answer for clarity. – emil Jan 15 '22 at 10:53
5

up to a certain point. I still get this error after uninstalling both `jax` and `jaxlib`, installing with `pip install --upgrade jax jaxlib` and then trying to use it as part of Trax: `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.` Same OS version and build, python 3.9.7 within a conda environment. – arturomp Jan 20 '22 at 22:14
@arturomp see my answer that I just posted. I fixed this on my machine by installing from conda-forge directly. – Nicholas G Reich Feb 08 '22 at 17:50
thank you! I think I did something similar: https://stackoverflow.com/a/70815865/ – arturomp Feb 08 '22 at 17:52

importing jax fails on mac with m1 chip

4 Answers4

Linked