CPU Support / Necessary AVX Instructions

[DerToast3r]

Hi,

the last week I struggled on using/installing jax.
I used a quite old CPU, an Intel Xeon x5670 and had trouble importing the jaxlib with cuda and cudnn support. It failed at importing and told me that my system is not compatible to the required AVX instructions.
Compiling it myself didnt solve the problem, so I put my GPU and Harddrive into an Intel 9th gen system, after setting the flag (os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '.7') to reduce preallocation on my display gpu and reinstalling jax&jaxlib version 0.3.14, it worked.

Now to the question, according to [https://en.wikipedia.org/wiki/Advanced_Vector_Extensions], can I assume that Intel Sandy Bridge architecture is the minimum requirement for jax with cuda support?

[hawkinsp]

JAX will work on CPUs without AVX instructions, but you must build jaxlib from source. The binary wheels we ship are built assuming at a minimum AVX support. This mostly affects folks with very old CPUs (e.g., your CPU is 12 years old) or using emulators like Rosetta that do not support AVX.

So my best suggestion to you is: build your own jaxlib wheels from source! See https://jax.readthedocs.io/en/latest/developer.html for details.