RuntimeError: This version of jaxlib was built using AVX instructions, which your CPU and/or operating system do not support.

[PaulScemama]

Hi,

I am trying to install JAX on a remote virtual machine that (as far as I know) runs Ubuntu and Linux x86_64. However, after installing and running

import jax

I get the error in the title:

RuntimeError: This version of jaxlib was built using AVX instructions, which your CPU and/or operating system do not support. 

Any ideas what could be going on on my side that would cause such an error? And if there is more specific information you need about the OS / CPU please let me know! Thanks

[jakevdp]

See #9580 – we don't provide pre-built wheels for non-AVX systems, but you could probably build jaxlib yourself (see https://github.com/google/jax/blob/main/CHANGELOG.md#jaxlib-0162-march-9-2021). Though I'm not aware of anyone building non-AVX jax in recent versions.

[jakevdp]

Also, this answer seems relevant if you want to check what features your processor enables: https://unix.stackexchange.com/questions/43539/what-do-the-flags-in-proc-cpuinfo-mean

[PaulScemama]

Thanks @jakevdp ! This helped a lot and I figured it out.

[hongzhouye]

@PaulScemama Hi, may I ask how you resolved your problem? I am using a 2021 macbook pro with a M1 chip and also see this error. I tried different versions of jax (the earliest in my pip is 0.4.6) but none of them seems to work. Thank you in advance for the help!

[PaulScemama]

hey @hongzhouye, my issue was I had accidentally wiped the AVX extensions, so installing them resolved my issue. However, if your machine does not support AVX extensions, you are out of luck unless you want to build jaxlib from source. @jakevdp links a page with some info to do so.

I also realize there have been some updates since I originally posted this, so maybe @jakevdp has some better advice. Have you tried installing the latest version here?

Sorry I can't be of more help.

[hongzhouye]

@PaulScemama Thank you for your quick response 🙏! After more searching and reading, I realized that my problem was that my Python was inherited from an older mac using intel chip. After installing an ARM64-Python from Miniforge, JAX now runs on my machine. I put the link to discussion here in case someone may find it useful later.

[PaulScemama]

@hongzhouye sounds good! You might be interested in using 'pyenv' to manage python versions in the future -- makes it very very easy. I personally use 'pyenv' for managing python versions and 'virtualenv' for creating environments with different package installs.

[hongzhouye]

@PaulScemama Thank you for the suggestions 🙏 !