Jax install error on Mac M1

[sjfleming]

This is not your bug, it's a Jax bug here:
jax-ml/jax#7237

BUT... it's preventing some people in our group from using scvi-tools versions since jax was made a requirement. Our work-issued laptops are Mac M1, and so we cannot pip install scvi-tools anymore due to the Jax requirement.

Is it possible to make Jax an install option? And not an actual requirement? We don't want to use Jax ourselves.

pip install scvi-tools

import scvi

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.

Versions:

all versions 0.15.x +

[adamgayoso]

On M1 mac, it's possible to install jax via conda (really jaxlib is the issue but you can see arm64 here), where there are compatible installations for M1. Would that be a workable solution for your team?

Secondarily, on M1 Mac it's important to use a native version of Python (arm64), which can be installed via the appropriate miniconda/miniforge installer.

[adamgayoso]

According to this issue jax-ml/jax#5501 there are m1 (arm64) wheels on pypi, so I would ensure native versions of python are being used

[sjfleming]

Thanks so much for the pointers @adamgayoso , it seems I need to look into this more. :)
I'll let you know!

[sjfleming]

Alright, so thank you so much for the thorough answer! We were in fact running an non-arm64 python install.

The correct python install, with a simple pip install of scvi-tools did work, for version 0.19.0

Someone on my team told me that the latest version 0.20.1 threw some sort of error related to jax, but unfortunately I don't have the details at this time. If I figure out that there is a real issue there, I will reopen this.

Thanks again!

[jesswhitts]

Hello,

I'm having this same issue - running anaconda3 on an M1 mac.

How did you resolve this?
Not sure how to tell if I'm using arm64 or not, or how to correct it.

[sjfleming]

Maybe try this link to install miniconda?

https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh

To figure out what you've got, type conda info

[jesswhitts]

Sorry for delayed response - but this did the trick. Thanks!