Installing jaxspec on Mac M1

renecotyfanboy / jaxspec

jaxspec is an X-ray spectra Bayesian analysis package, relying on JAX to enable just in time compilation

https://jaxspec.readthedocs.io/en/latest/

MIT License

23 stars 1 forks source link

Installing jaxspec on Mac M1 #131

Open CamilleDiez opened 7 months ago

CamilleDiez commented 7 months ago

At this time installation on mac with Apple silicon is not working with pip. There is a workaround which consists in uninstalling jax and jaxlib and reinstalling with conda which comes with its own wheels

pip uninstall jax jaxlib
conda install -c conda-forge jaxlib
conda install -c conda-forge jax

It just works. LF updates with jax-metal to get GPU integration

CamilleDiez commented 7 months ago

Code fails when compiled with jit and installing jax-metal with pip does not work here