Problems with jax in mac & windows

[Lucas-Maciel]

Hi,

I have tried to run cell2fate in my Mac M2, however, I can't due to problems with jaxlib.

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[1], [line 1](vscode-notebook-cell:?execution_count=1&line=1)
----> [1](vscode-notebook-cell:?execution_count=1&line=1) import cell2fate as c[2](vscode-notebook-cell:?execution_count=1&line=2)f
      2 import scanpy as sc
      [3](vscode-notebook-cell:?execution_count=1&line=3) import numpy as np

File [~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/cell2fate/__init__.py:11](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/cell2fate/__init__.py:11)
      [8](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/cell2fate/__init__.py:8) from pyro.distributions.transforms import SoftplusTransform
      [9](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/cell2fate/__init__.py:9) from torch.distributions import biject_to, transform_to
---> [11](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/cell2fate/__init__.py:11) from ._cell2fate_DynamicalModel import \
     [12](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/cell2fate/__init__.py:12) Cell2fate_DynamicalModel
     [15](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/cell2fate/__init__.py:15) try:
     [16](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/cell2fate/__init__.py:16)     import importlib.metadata as importlib_metadata

File [~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/cell2fate/_cell2fate_DynamicalModel.py:9](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/cell2fate/_cell2fate_DynamicalModel.py:9)
      [7](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/cell2fate/_cell2fate_DynamicalModel.py:7) from anndata import AnnData
      [8](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/cell2fate/_cell2fate_DynamicalModel.py:8) from pyro import clear_param_store
----> [9](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/cell2fate/_cell2fate_DynamicalModel.py:9) from scvi.model._utils import parse_use_gpu_arg
     [10](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/cell2fate/_cell2fate_DynamicalModel.py:10) from scvi.dataloaders import AnnDataLoader
     [11](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/cell2fate/_cell2fate_DynamicalModel.py:11) from scvi.utils import track

File [~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/__init__.py:10](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/__init__.py:10)
      [7](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/__init__.py:7) from ._settings import settings
      [9](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/__init__.py:9) # this import needs to come after prior imports to prevent circular import
---> [10](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/__init__.py:10) from . import data, model, external, utils
     [12](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/__init__.py:12) # https://github.com/python-poetry/poetry/pull/2366#issuecomment-652418094
     [13](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/__init__.py:13) # https://github.com/python-poetry/poetry/issues/144#issuecomment-623927302
     [14](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/__init__.py:14) try:

File [~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/__init__.py:24](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/__init__.py:24)
      [1](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/__init__.py:1) from anndata import read_csv, read_h5ad, read_loom, read_text
      [3](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/__init__.py:3) from ._datasets import (
      [4](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/__init__.py:4)     annotation_simulation,
      [5](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/__init__.py:5)     brainlarge_dataset,
   (...)
     [22](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/__init__.py:22)     synthetic_iid,
     [23](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/__init__.py:23) )
---> [24](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/__init__.py:24) from ._manager import AnnDataManager
     [25](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/__init__.py:25) from ._preprocessing import (
     [26](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/__init__.py:26)     organize_cite_seq_10x,
     [27](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/__init__.py:27)     organize_multiome_anndatas,
     [28](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/__init__.py:28)     poisson_gene_selection,
     [29](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/__init__.py:29) )
     [30](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/__init__.py:30) from ._read import read_10x_atac, read_10x_multiome

File [~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/_manager.py:15](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/_manager.py:15)
     [12](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/_manager.py:12) from anndata import AnnData
     [14](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/_manager.py:14) import scvi
---> [15](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/_manager.py:15) from scvi._types import AnnDataField
     [16](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/_manager.py:16) from scvi.utils import attrdict
     [18](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/data/_manager.py:18) from . import _constants

File [~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/_types.py:3](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/_types.py:3)
      [1](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/_types.py:1) from typing import Dict, Type, Union
----> [3](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/_types.py:3) import jax.numpy as jnp
      [4](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/_types.py:4) import torch
      [6](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/scvi/_types.py:6) from scvi.data.fields import BaseAnnDataField

File [~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/__init__.py:35](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/__init__.py:35)
     [30](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/__init__.py:30) del _cloud_tpu_init
     [32](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/__init__.py:32) # Confusingly there are two things named "config": the module and the class.
     [33](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/__init__.py:33) # We want the exported object to be the class, so we first import the module
     [34](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/__init__.py:34) # to make sure a later import doesn't overwrite the class.
---> [35](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/__init__.py:35) from jax import config as _config_module
     [36](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/__init__.py:36) del _config_module
     [38](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/__init__.py:38) # Force early import, allowing use of `jax.core` after importing `jax`.

File [~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/config.py:17](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/config.py:17)
      [1](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/config.py:1) # Copyright 2018 The JAX Authors.
      [2](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/config.py:2) #
      [3](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/config.py:3) # Licensed under the Apache License, Version 2.0 (the "License");
   (...)
     [14](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/config.py:14) 
     [15](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/config.py:15) # TODO(phawkins): fix users of this alias and delete this file.
---> [17](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/config.py:17) from jax._src.config import config  # noqa: F401

File [~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/_src/config.py:24](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/_src/config.py:24)
     [21](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/_src/config.py:21) import threading
     [22](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/_src/config.py:22) from typing import Any, List, Callable, Hashable, NamedTuple, Iterator, Optional
---> [24](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/_src/config.py:24) from jax._src import lib
     [25](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/_src/config.py:25) from jax._src.lib import jax_jit
     [26](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/_src/config.py:26) from jax._src.lib import transfer_guard_lib

File [~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/_src/lib/__init__.py:84](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/_src/lib/__init__.py:84)
     [80](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/_src/lib/__init__.py:80) # Before importing any C compiled modules from jaxlib, first import the CPU
     [81](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/_src/lib/__init__.py:81) # feature guard module to verify that jaxlib was compiled in a way that only
     [82](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/_src/lib/__init__.py:82) # uses instructions that are present on this machine.
     [83](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/_src/lib/__init__.py:83) import jaxlib.cpu_feature_guard as cpu_feature_guard
---> [84](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/_src/lib/__init__.py:84) cpu_feature_guard.check_cpu_features()
     [86](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/_src/lib/__init__.py:86) # TODO(phawkins): remove after minimium jaxlib version is 0.4.9 or newer.
     [87](~/miniconda3/envs/cell2fate_env/lib/python3.9/site-packages/jax/_src/lib/__init__.py:87) try:

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.

I have tried to install using the whl file proper for mac, however, I still had the same problem.

pip install -U ~/Downloads/jaxlib-0.4.10-cp39-cp39-macosx_10_14_x86_64.whl

I also tried to install on one Windows machine and it also had problems with jax (for other packages in the past I had to use the legacy version of jax).

Do you have any recommendation?

Thanks

[AlexanderAivazidis]

Hi Lucas! We just tried to install cell2fate on an M2 mac to see if we get similar errors, but it works fine for us. However, it is very slow because it does not use the GPU, even after some tweaks. So overall we recommend not to use it on a Mac at this point. We plan to make a docker image soon, so people can use cell2fate on any operating system.
In the meantime, the recommendation would be to use windows and to use the installation instructions here: https://cell2fate.readthedocs.io/en/latest/installation.html
Which includes making a fully seperate conda environment. Have you done this? I think a fully seperate conda environment may solve your problems. You may also have to tell python to NOT use user site for installing packages as explained at the very end of the installation instructions.

[lcw-12323]

@AlexanderAivazidis Hello, I have a similar problem when deploying cell2fate using windows. I can't find jaxlib-0.4.10 for windows, the smallest version available is 0.4.13, how do I get the right version to run cell2fate on windows?