Jax on MacOS ARM (M1)

[amir-f]

Description

I'm new to Jax, have a MacOS ARM and the initial setup wasn't quite smooth.

pip install -U pip
pip install jax
pip install jaxlib

And then I get

>>> import jax
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/amir/.venv/lib/python3.8/site-packages/jax/__init__.py", line 35, in <module>
    from jax import config as _config_module
  File "/Users/amir/.venv/lib/python3.8/site-packages/jax/config.py", line 17, in <module>
    from jax._src.config import config  # noqa: F401
  File "/Users/amir/.venv/lib/python3.8/site-packages/jax/_src/config.py", line 28, in <module>
    from jax._src import lib
  File "/Users/amir/.venv/lib/python3.8/site-packages/jax/_src/lib/__init__.py", line 85, 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 thought the support was there. Is the recommendation to build from source? I want to run it locally on the laptop for initial prototyping.

What jax/jaxlib version are you using?

jax 0.4.1, jaxlib 0.4.1

Which accelerator(s) are you using?

CPU: Apple M1 Max

Additional system info

Python 3.8.2, Mac OS 13.1 (22C65)

NVIDIA GPU info

No response

[apaszke]

Could you please share the installation logs? I just checked the PyPI page for jaxlib and you can see macOS ARM packages in there.

[mlsw]

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.

This suggests that you are using a x86_64 version of Python under Rosetta, rather than a native arm64 version. You can verify this by running:

python -c 'import subprocess; subprocess.run(["sysctl", "-n", "sysctl.proc_translated"])'

A value of 0 is returned for a native process, while 1 will be returned if Rosetta is being used.

[amir-f]

@mlsw very interesting. It does return 1. Let me look into how to fix my setup then.

@apaszke I looked at the pip installation logs pip install jaxlib --log /tmp/piplog.txt and it indeed is using the x86 wheel.

2023-01-11T03:05:45,869   Using cached jaxlib-0.4.1-cp38-cp38-macosx_10_14_x86_64.whl (64.0 MB)
2023-01-11T03:05:46,121   Added jaxlib from https://files.pythonhosted.org/packages/3a/29/a897b31a798c5e04d03d1f262d8028eac92b392b57136be3c5ad913630be/jaxlib-0.4.1-cp38-cp38-macosx_10_14_x86_64.whl to build tracker '/private/var/folders/bz/mhjxkd6s72n_95lc782sx0240000gq/T/pip-build-tracker-4sw69uh6'

[amir-f]

Fixed it by replacing x86 python with the arm version (brew install helped). Perhaps the cpu_feature_guard.check_cpu_features() code can detect whether it's under Rosetta and display a better error message. Closing this for now.

[mjjohns1]

I'm encountering this issue when trying to build Jax from source. It's trying to build with cpu=x86_64 but should be arm64 version. @amir-f can you provide more details about how you replaced your x86 python with the arm version?

[amir-f]

@mjjohns1 I did not build from source. I first uninstalled the existing Python installation. See here for more info.

Then installed the M1 version via brew install python. Brew picks ARM binaries (see here) on Apple M1 machines.

If you are building jax from source I believe you need to set the target platform for bazel. Among the configs I only see one macos config. If bazel --config macos doesn't target ARM I'm not sure what the correct flags are.