Cannot import JAX using M1 Mac

Questions v5
1

Hi!

I have just installed version 4 of PyMC on an Apple m1 device and I followed these instructions.

However, when I try to import jax, I get the following error:

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.

How can I fix this?

P.S. I have recently started working with PyMC and I am learning a lot. I cannot wait to use this new version. Thank you for your hard work!

6 Likes
2

I running into this issue as well. Did you figure it out eventually?

1 Like
3

Note that the “official” installation instructions can be found here. What version of jax do you have installed and how was it installed?

4

Installed jax and jaxlib == 0.3.0 using pip install.

5

I have M1, and I don’t have this message. I used:

python3 -m pip install jax jaxlib numpyro

6

Good to hear! What versions of jax, jaxlib, numpyro and python are you on?

7 
>>> print(watermark(packages="jax,jaxlib,numpyro"))
jax    : 0.3.23
jaxlib : 0.3.22
numpyro: 0.10.1

# no messages
>>> import jaxlib
>>> import jax

Python is 3.10.8

8

Finally got it to work! I re-installed mini forge, this time not using homebrew. (Ad modem Mac M1 Monterey Installing Miniforge and Anaconda/Miniconda Side-by-Side - YouTube).

Hang in there, j_catulo, the sweet sweet performance increase is worth it.

3 Likes
9

Ah yes. I definitely suggest using the recommended installation routine at all time, regardless of platform. Glad it worked out for you!

10

I used homebrew in my case - so, maybe there was something else