Issues installing jax (conda, SL7)

[kmnelson]

I am attempting to install jax in a conda environment. My machine is running Scientific Linux 7.9

I ran:
conda create -n test python=3
To create a conda environment and then
my/path/anaconda3/envs/test/bin/pip install jax jaxlib
Which seems fine, but if I then run
my/path/anaconda3/envs/test/bin/python
import jax
I get a printout that only says "Illegal instruction" (importing jaxlib does not cause any errors)

I have also tried unsuccessfully building from source.

Does anyone have tips for installing jax? Help is greatly appreciated

-Kevin

[jakevdp]

Does your CPU support AVX? The pre-built JAX binaries for version 0.1.62 and newer require AVX features; you could try installing an older version of JAX to see if that fixes the issue:

$ pip install jax==0.2.9 jaxlib==0.1.61

If that does fix the issue, then it's likely that your hardware does not support AVX. If that is the case, then the only way to use newer jaxlib on your hardware would be to compile jaxlib from source. See Building JAX From Source for information about that, and don't forget the extra build flags mentioned in the changelog link above.

[kmnelson]

Yes, installing the older version did work. I will try building jaxlib from source with the appropriate flags.

Thanks!!