Cannot import JAX using M1 Mac


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!