JAX error when importing jax.numpy under M2

[coolalu]

Hi there,

I was trying to run LMMM package on my new M2 laptop, but on the very beginning code where it tried to import jax.numpy, it returned an error saying the version of jaxlib is not supported under current CPU/or operating system. Has anyone experienced the same issue?

Attached screenshots of code and error message for reference:

Thanks!

[BigDataRPG]

Hi @coolalu

I once got problem also in M2. Is it possible reinstall respectively

pip install numpy cython pip install jax jaxlib numpyro

If it does not work, you can try re install in new python environment. Try install --no-deps

` pip install --no-deps lightweight_mmm pip install numpy cython pip install jax jaxlib numpyro

No dependency module when run lightweight_mmm pip install --no-deps absl-py pip install --no-deps immutabledict pip install --no-deps statsmodels pip install --no-deps patsy

ModuleNotFoundError: No module named 'tensorflow' pip install tensorflow-macos pip install tensorflow-metal

Another dependency module when run lightweight_mmm pip install --no-deps arviz pip install --no-deps xarray pip install --no-deps netCDF4 pip install --no-deps cftime pip install --no-deps xarray-einstats `

Hope it would helps you

[coolalu]

Hi @BigDataRPG,

Thank you for sharing! I tried to reinstall them and the error went away! However, a new attribution error popped up -

Have you experienced the same?

I found there is a post discussing it https://github.com/google/jax/discussions/14036, but it didn't resolve the issue.

Thanks!