Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Cannot install recent version of Jax/Jaxlib on MacOS Sonoma #19067

Closed
fonnesbeck opened this issue Dec 20, 2023 · 22 comments
Closed

Cannot install recent version of Jax/Jaxlib on MacOS Sonoma #19067

fonnesbeck opened this issue Dec 20, 2023 · 22 comments
Assignees
@jakevdp
Labels
bug Something isn't working

Comments

@fonnesbeck
Copy link

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
@fonnesbeck fonnesbeck added the bug Something isn't working label Dec 20, 2023
@jakevdp
Copy link
Collaborator

jakevdp commented Dec 20, 2023 •
edited

Hi Chris, thanks for the report! This is definitely not intended. What OSX version are you using?

Also, what happens if you attempt to install one of the osx ARM wheels directly? e.g.

$ pip install https://storage.googleapis.com/jax-releases/max/jaxlib-0.4.23-cp311-cp311-macosx_11_0_arm64.whl

or for OSX 10.14:

$ pip install https://storage.googleapis.com/jax-releases/mac/jaxlib-0.4.23-cp311-cp311-macosx_10_14_x86_64.whl

The full list of available pre-built wheels can be found at https://storage.googleapis.com/jax-releases/jax_releases.html

@jakevdp
Copy link
Collaborator

jakevdp commented Dec 20, 2023 •
edited

Ah, just noticed that you mentioned Catalina. It looks like Catalina wheels are not uploaded to PyPI (looking at the list here), but they are in the cloud storage bucket. So you could probably install JAX this way, assuming the 10_14 wheel there is the right one for your system

$ pip install --upgrade "jax[cpu]" -f https://storage.googleapis.com/jax-releases/jax_releases.html

@fonnesbeck fonnesbeck changed the title Cannot install recent version of Jax/Jaxlib on MacOS Catalina Cannot install recent version of Jax/Jaxlib on MacOS Sonoma Dec 20, 2023
@fonnesbeck
Copy link
Author

Hey Jake! My bad there, I wrote Catalina when I meant Sonoma (14.1).

@fonnesbeck
Copy link
Author

fonnesbeck commented Dec 20, 2023 •
edited

Since I've now established that I'm on 14.1, I went ahead and tried the direct wheel install. It fails as follows, even though the OS version is correct:

❯ pip install https://storage.googleapis.com/jax-releases/mac/jaxlib-0.4.23-cp311-cp311-macosx_10_14_x86_64.whl
ERROR: jaxlib-0.4.23-cp311-cp311-macosx_10_14_x86_64.whl is not a supported wheel on this platform.
❯ sw_vers
ProductName:		macOS
ProductVersion:		14.1
BuildVersion:		23B74

Similar result for the ARM wheel.

@jakevdp
Copy link
Collaborator

jakevdp commented Dec 20, 2023 •
edited

10_14 is the Catalina wheel – did you try with the 11_0 wheel?

I just tried locally on my mac (running OSX Sonoma 14.2 and Python 3.9) and the jaxlib-0.4.23-cp39-cp39-macosx_11_0_arm64.whl worked fine.

@jakevdp
Copy link
Collaborator

jakevdp commented Dec 20, 2023 •
edited

The fact that pip ended up installing jaxlib 0.2.X suggests to me that you may be using a pip executable associated with an older Python version. pip with Python 3.11 shouldn't be able to install jax 0.2.10, because the Python version does match the requirements specified in that version's metadata.

If this is the problem on your system, then python -m pip install "jax[cpu]" might be the fix you need.

@fonnesbeck
Copy link
Author

If this is the problem on your system, then python -m pip install "jax[cpu]" might be the fix you need.

This replicates my original experience: it tries every version of jax, until it happily installs 2.10.0:

WARNING: jax 0.2.10 does not provide the extra 'cpu'
Installing collected packages: jax
Successfully installed jax-0.2.10

This is Python 3.11.5

@fonnesbeck
Copy link
Author

10_14 is the Catalina wheel – did you try with the 11_0 wheel?

I just tried locally on my mac (running OSX Sonoma 14.2 and Python 3.9) and the jaxlib-0.4.23-cp39-cp39-macosx_11_0_arm64.whl worked fine.

So jaxlib-0.4.23-cp311-cp311-macosx_11_0_arm64.whl should be the one I need then, right? This also errors out:

ERROR: jaxlib-0.4.23-cp311-cp311-macosx_11_0_arm64.whl is not a supported wheel on this platform.

@jakevdp
Copy link
Collaborator

jakevdp commented Dec 20, 2023 •
edited

Ah, jax[cpu] was added in 0.2.11, so if pip is not finding a compatible CPU jaxlib, it will keep dialing back until 0.2.10, which doesn't specify a jaxlib requirement.

I'm stumped as to why pip is treating that jaxlib wheel as being incompatible with your platform. That's the one that should work for you, if you're on OSX Sonoma with Python 3.11.

@jakevdp
Copy link
Collaborator

jakevdp commented Dec 20, 2023

What version of pip do you have? An outdated pip version may be incompatible with these more recent wheels.

@fonnesbeck
Copy link
Author

23.3.2. I did a pip upgrade just to be sure.

@jakevdp
Copy link
Collaborator

jakevdp commented Dec 20, 2023

Well, like I said, I'm stumped. Unfortunately pip doesn't offer much visibilty into why it considers any particular wheel file incompatible with the current system, so all we can do is guess.

If you do Apple-> About This Mac, what does it say for the Chip?

@fonnesbeck
Copy link
Author

fonnesbeck commented Dec 20, 2023 •
edited

OK, cool. There must be something up with my system. Thanks for checking!

(It's an M1 chip, to answer your question)

@jakevdp
Copy link
Collaborator

jakevdp commented Dec 20, 2023

To help us get on the same page, can you paste the output of these commands?

$ python --version
$ uname -a  # feel free to redact the hostname if you wish
$ sw_vers
$ python -m pip --version
$ python -m pip install "jax[cpu]" -U -v -f https://storage.googleapis.com/jax-releases/jax_releases.html

@fonnesbeck
Copy link
Author

❯ python --version
Python 3.11.5
❯ uname -a # feel free to redact the hostname if you wish
Darwin Bake.local 23.1.0 Darwin Kernel Version 23.1.0: Mon Oct 9 21:28:12 PDT 2023; root:xnu-10002.41.9~6/RELEASE_ARM64_T8103 arm64
❯ sw_vers
ProductName: macOS
ProductVersion: 14.1
BuildVersion: 23B74
❯ python -m pip --version
pip 23.3.2 from /Users/cfonnesbeck/mambaforge/envs/pymc/lib/python3.11/site-packages/pip (python 3.11)

And then the install attempt ends as before:

WARNING: jax 0.2.10 does not provide the extra 'cpu'
Installing collected packages: jax
Successfully installed jax-0.2.10

@jakevdp
Copy link
Collaborator

jakevdp commented Dec 20, 2023

What about this?

$ python -m pip debug --verbose | grep cp311

@jakevdp
Copy link
Collaborator

jakevdp commented Dec 20, 2023 •
edited

Ah, I think you might be hitting this issue: microsoft/onnxruntime#17166 (comment). 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?

@fonnesbeck
Copy link
Author

So, the pip debug returned nothing prior to the workaround you suggested, but after installing onyxruntime, it yields this list:

  cp311-cp311-macosx_10_16_arm64
  cp311-cp311-macosx_10_16_universal2
  cp311-cp311-macosx_10_15_arm64
  cp311-cp311-macosx_10_15_universal2
  cp311-cp311-macosx_10_14_arm64
  cp311-cp311-macosx_10_14_universal2
  cp311-cp311-macosx_10_13_arm64
  cp311-cp311-macosx_10_13_universal2
  cp311-cp311-macosx_10_12_arm64
  cp311-cp311-macosx_10_12_universal2
  cp311-cp311-macosx_10_11_arm64
  cp311-cp311-macosx_10_11_universal2
  cp311-cp311-macosx_10_10_arm64
  cp311-cp311-macosx_10_10_universal2
  cp311-cp311-macosx_10_9_arm64
  cp311-cp311-macosx_10_9_universal2
  cp311-cp311-macosx_10_8_arm64
  cp311-cp311-macosx_10_8_universal2
  cp311-cp311-macosx_10_7_arm64
  cp311-cp311-macosx_10_7_universal2
  cp311-cp311-macosx_10_6_arm64
  cp311-cp311-macosx_10_6_universal2
  cp311-cp311-macosx_10_5_arm64
  cp311-cp311-macosx_10_5_universal2
  cp311-cp311-macosx_10_4_arm64
  cp311-cp311-macosx_10_4_universal2
  cp311-cp311-macosx_10_3_arm64
  cp311-cp311-macosx_10_3_universal2
  cp311-cp311-macosx_10_2_arm64
  cp311-cp311-macosx_10_2_universal2
  cp311-cp311-macosx_10_1_arm64
  cp311-cp311-macosx_10_1_universal2
  cp311-cp311-macosx_10_0_arm64
  cp311-cp311-macosx_10_0_universal2
  cp311-abi3-macosx_10_16_arm64
  cp311-abi3-macosx_10_16_universal2
  cp311-abi3-macosx_10_15_arm64
  cp311-abi3-macosx_10_15_universal2
  cp311-abi3-macosx_10_14_arm64
  cp311-abi3-macosx_10_14_universal2
  cp311-abi3-macosx_10_13_arm64
  cp311-abi3-macosx_10_13_universal2
  cp311-abi3-macosx_10_12_arm64
  cp311-abi3-macosx_10_12_universal2
  cp311-abi3-macosx_10_11_arm64
  cp311-abi3-macosx_10_11_universal2
  cp311-abi3-macosx_10_10_arm64
  cp311-abi3-macosx_10_10_universal2
  cp311-abi3-macosx_10_9_arm64
  cp311-abi3-macosx_10_9_universal2
  cp311-abi3-macosx_10_8_arm64
  cp311-abi3-macosx_10_8_universal2
  cp311-abi3-macosx_10_7_arm64
  cp311-abi3-macosx_10_7_universal2
  cp311-abi3-macosx_10_6_arm64
  cp311-abi3-macosx_10_6_universal2
  cp311-abi3-macosx_10_5_arm64
  cp311-abi3-macosx_10_5_universal2
  cp311-abi3-macosx_10_4_arm64
  cp311-abi3-macosx_10_4_universal2
  cp311-abi3-macosx_10_3_arm64
  cp311-abi3-macosx_10_3_universal2
  cp311-abi3-macosx_10_2_arm64
  cp311-abi3-macosx_10_2_universal2
  cp311-abi3-macosx_10_1_arm64
  cp311-abi3-macosx_10_1_universal2
  cp311-abi3-macosx_10_0_arm64
  cp311-abi3-macosx_10_0_universal2
  cp311-none-macosx_10_16_arm64
  cp311-none-macosx_10_16_universal2
  cp311-none-macosx_10_15_arm64
  cp311-none-macosx_10_15_universal2
  cp311-none-macosx_10_14_arm64
  cp311-none-macosx_10_14_universal2
  cp311-none-macosx_10_13_arm64
  cp311-none-macosx_10_13_universal2
  cp311-none-macosx_10_12_arm64
  cp311-none-macosx_10_12_universal2
  cp311-none-macosx_10_11_arm64
  cp311-none-macosx_10_11_universal2
  cp311-none-macosx_10_10_arm64
  cp311-none-macosx_10_10_universal2
  cp311-none-macosx_10_9_arm64
  cp311-none-macosx_10_9_universal2
  cp311-none-macosx_10_8_arm64
  cp311-none-macosx_10_8_universal2
  cp311-none-macosx_10_7_arm64
  cp311-none-macosx_10_7_universal2
  cp311-none-macosx_10_6_arm64
  cp311-none-macosx_10_6_universal2
  cp311-none-macosx_10_5_arm64
  cp311-none-macosx_10_5_universal2
  cp311-none-macosx_10_4_arm64
  cp311-none-macosx_10_4_universal2
  cp311-none-macosx_10_3_arm64
  cp311-none-macosx_10_3_universal2
  cp311-none-macosx_10_2_arm64
  cp311-none-macosx_10_2_universal2
  cp311-none-macosx_10_1_arm64
  cp311-none-macosx_10_1_universal2
  cp311-none-macosx_10_0_arm64
  cp311-none-macosx_10_0_universal2
  cp311-none-any

However, the workaround does not result in a good install -- it still settles for 0.2.10.

Interestingly, if I simply install via python -m pip install jax (without [cpu]) it will happily install jax-0.4.23, BUT is subsequently unable to install an associated jaxlib: ERROR: No matching distribution found for jaxlib

@jakevdp
Copy link
Collaborator

jakevdp commented Dec 20, 2023 •
edited

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.

@jakevdp
Copy link
Collaborator

jakevdp commented Dec 20, 2023 •
edited

Interestingly, if I simply install via python -m pip install jax (without [cpu]) it will happily install jax-0.4.23,

This is because jax is a universal wheel, which your pip recognizes. jaxlib is an arm64 wheel, which your pip does not recognize.

@fonnesbeck
Copy link
Author

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!

@jakevdp
Copy link
Collaborator

jakevdp commented Dec 20, 2023 •
edited

No worries, this was a very tricky one. Glad it's resolved! Hope we can catch up IRL sometime 😀

@jakevdp jakevdp self-assigned this Dec 20, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@jakevdp jakevdp

Labels
bug Something isn't working
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
2 participants
@fonnesbeck @jakevdp