Issue with installing jax/jaxlib for Mac OSX (Sonoma v14.5) with M1 chip

[hongzhouye]

Hi community,

I am new to jax and hope to get some insights into the issue I encountered with installing jax on my Mac OSX machine with M1 chip. I used the following shell script for installing and testing jax.

#!/bin/bash

`which python3.11` -m venv jax_env
source jax_env/bin/activate
pip install jax jaxlib
python -c "import jax"

It generated the following screen output:

Collecting jax
  Obtaining dependency information for jax from https://files.pythonhosted.org/packages/fd/f2/9dbb75de3058acfd1600cf0839bcce7ea391148c9d2b4fa5f5666e66f09e/jax-0.4.30-py3-none-any.whl.metadata
  Using cached jax-0.4.30-py3-none-any.whl.metadata (22 kB)
Collecting jaxlib
  Obtaining dependency information for jaxlib from https://files.pythonhosted.org/packages/33/2d/b6078f5d173d3087d32b1b49e5f65d406985fb3894ff1d21905972b9c89d/jaxlib-0.4.30-cp311-cp311-macosx_10_14_x86_64.whl.metadata
  Downloading jaxlib-0.4.30-cp311-cp311-macosx_10_14_x86_64.whl.metadata (1.0 kB)
Collecting ml-dtypes>=0.2.0 (from jax)
  Obtaining dependency information for ml-dtypes>=0.2.0 from https://files.pythonhosted.org/packages/42/6b/b2fa3e2386c2b7dde43f12b83c67f6e583039141dfbb58e5c8fd365a5a7d/ml_dtypes-0.4.0-cp311-cp311-macosx_10_9_universal2.whl.metadata
  Downloading ml_dtypes-0.4.0-cp311-cp311-macosx_10_9_universal2.whl.metadata (20 kB)
Collecting numpy>=1.22 (from jax)
  Obtaining dependency information for numpy>=1.22 from https://files.pythonhosted.org/packages/4f/c1/42d1789f1dff7b65f2d3237eb88db258a5a7fdfb981b895509887c92838d/numpy-2.0.0-cp311-cp311-macosx_14_0_x86_64.whl.metadata
  Downloading numpy-2.0.0-cp311-cp311-macosx_14_0_x86_64.whl.metadata (60 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 60.9/60.9 kB 2.2 MB/s eta 0:00:00
Collecting opt-einsum (from jax)
  Obtaining dependency information for opt-einsum from https://files.pythonhosted.org/packages/bc/19/404708a7e54ad2798907210462fd950c3442ea51acc8790f3da48d2bee8b/opt_einsum-3.3.0-py3-none-any.whl.metadata
  Using cached opt_einsum-3.3.0-py3-none-any.whl.metadata (6.5 kB)
Collecting scipy>=1.9 (from jax)
  Obtaining dependency information for scipy>=1.9 from https://files.pythonhosted.org/packages/6a/d6/db686519059afb367e5a06935556b50fa422d792a658ce071f4527c785bf/scipy-1.14.0-cp311-cp311-macosx_14_0_x86_64.whl.metadata
  Downloading scipy-1.14.0-cp311-cp311-macosx_14_0_x86_64.whl.metadata (60 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 60.8/60.8 kB 5.8 MB/s eta 0:00:00
Using cached jax-0.4.30-py3-none-any.whl (2.0 MB)
Downloading jaxlib-0.4.30-cp311-cp311-macosx_10_14_x86_64.whl (83.5 MB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 83.5/83.5 MB 4.9 MB/s eta 0:00:00
Downloading ml_dtypes-0.4.0-cp311-cp311-macosx_10_9_universal2.whl (390 kB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 390.9/390.9 kB 5.1 MB/s eta 0:00:00
Downloading numpy-2.0.0-cp311-cp311-macosx_14_0_x86_64.whl (6.9 MB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 6.9/6.9 MB 3.2 MB/s eta 0:00:00
Downloading scipy-1.14.0-cp311-cp311-macosx_14_0_x86_64.whl (25.5 MB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 25.5/25.5 MB 2.3 MB/s eta 0:00:00
Using cached opt_einsum-3.3.0-py3-none-any.whl (65 kB)
Installing collected packages: numpy, scipy, opt-einsum, ml-dtypes, jaxlib, jax
Successfully installed jax-0.4.30 jaxlib-0.4.30 ml-dtypes-0.4.0 numpy-2.0.0 opt-einsum-3.3.0 scipy-1.14.0

[notice] A new release of pip is available: 23.2.1 -> 24.1.1
[notice] To update, run: pip install --upgrade pip
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/Users/hzye/local/opt/jax/jax_env/lib/python3.11/site-packages/jax/__init__.py", line 25, in <module>
    from jax._src.cloud_tpu_init import cloud_tpu_init as _cloud_tpu_init
  File "/Users/hzye/local/opt/jax/jax_env/lib/python3.11/site-packages/jax/_src/cloud_tpu_init.py", line 17, in <module>
    from jax._src import config
  File "/Users/hzye/local/opt/jax/jax_env/lib/python3.11/site-packages/jax/_src/config.py", line 27, in <module>
    from jax._src import lib
  File "/Users/hzye/local/opt/jax/jax_env/lib/python3.11/site-packages/jax/_src/lib/__init__.py", line 84, 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. This error is frequently encountered on macOS when running an x86 Python installation on ARM hardware. In this case, try installing an ARM build of Python. Otherwise, you may be able work around this issue by building jaxlib from source.

One can see that the installation was successful but an error occurred when the last line import jax was executed. The problem seems to be using x86 Python installation on ARM hardware, but I wasn't so sure.

Has anyone else shared similar experience? And if so, how did you resolve the issue? Many thanks in advance 🙏 !

(I found several discussions about the incompatibility between jax and M1 chips before 2022. But to the best of my knowledge, this problem was resolved and jax now officially supports M1 chips.)

[dfm]

The error message is correct here! As you can see from the installation logs, all of the downloaded wheels have the x86_64 suffix in their filenames. That means that you're running Python in x86 emulator mode on your machine. To use JAX, you'll need to install an ARM build of Python, and as an added benefit all of your Python code will run faster!

I usually install Python using the Miniforge installer (make sure you choose the arm64 version!), but there are other options.

[hongzhouye]

Thank you @dfm! I will try installing ARM-built Python through Miniforge.

[hongzhouye]

When following the instruction on the miniforge github page for installation, I realized that command like this

curl -L -O "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh"

may still download the x86_64 version because uname -m returns x86_64 even on an M1 mac (this may be the deep reason causing the original issue I saw in the first place). This same issue has been reported by others as well, see e.g., here and here.

The solution I found in this context is replacing $(uname -m) with arm64 explicitly, i.e.,

curl -L -O "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-arm64.sh"

Then installing jax with this arm64 python resolves the original issue I posted.

@dfm thank you again for pointing me to miniforge!

[hawkinsp]

@hongzhouye Interesting. uname -m returns "arm64" for me on a Mac M1. My guess is that perhaps your shell is running in x86 emulation also!