jax: Cannot install recent version of Jax/Jaxlib on MacOS Sonoma

Description

The instructions for installing Jax on a Mac ARM machine on the Jax website are as follows:

pip install --upgrade pip
pip install --upgrade "jax[cpu]"

However, when I attempt to run this, pip attempts dozens of older versions, and finally installs version 0.2.10. Is there no way to get version 0.4.x to run on a Mac?

What jax/jaxlib version are you using?

NA

Which accelerator(s) are you using?

CPU

Additional system info?

Python 3.11

NVIDIA GPU info

No response

About this issue

  • Original URL
  • State: closed
  • Created 6 months ago
  • Comments: 22

Most upvoted comments

Ah, got it. Yes, can confirm that SYSTEM_VERSION_COMPAT=0 pip install jax[cpu] does the trick! Sorry for being a bit dense here – its been a long day. Thanks for the help!

Thanks - it looks like pip is not recognizing that 11_0 wheels are compatible with your system. That’s similar to the issue I linked above, and probably related to the SDK version used to build the Python you are using.

I don’t think onyxruntime is relevant here, no need to install it – I was just linking to that thread because it looks like they are discussing a similar issue.

Ah, I think you might be hitting this issue: https://github.com/microsoft/onnxruntime/issues/17166#issuecomment-1693938597. It seems like if pip is build with an older OSX SDK it may not recognize some tags (the pip debug command above would help show whether that’s the case).

Maybe try the workaround mentioned in that comment?