Protocol unirep: ImportError: cannot import name 'cusolver' from 'jaxlib' (/mnt/lsf-nas-1/os-shared/anaconda3/envs/hannes_bio_embeddings/lib/python3.7/site-packages/jaxlib/__init__.py)

[HannesStark]

Metadata

key	value
version	0.2.1
cuda	True

Parameter

key	value
type	embed
protocol	unirep
reduce	True

Traceback

Traceback (most recent call last):
  File "/mnt/lsf-nas-1/os-shared/anaconda3/envs/hannes_bio_embeddings/lib/python3.7/site-packages/bio_embeddings/utilities/pipeline.py", line 284, in execute_pipeline_from_config
    stage_output_parameters = stage_runnable(**stage_parameters)
  File "/mnt/lsf-nas-1/os-shared/anaconda3/envs/hannes_bio_embeddings/lib/python3.7/site-packages/bio_embeddings/embed/pipeline.py", line 400, in run
    embedder: EmbedderInterface = embedder_class(**result_kwargs)
  File "/mnt/lsf-nas-1/os-shared/anaconda3/envs/hannes_bio_embeddings/lib/python3.7/site-packages/bio_embeddings/embed/unirep_embedder.py", line 33, in __init__
    from jax_unirep.utils import load_params_1900
  File "/mnt/lsf-nas-1/os-shared/anaconda3/envs/hannes_bio_embeddings/lib/python3.7/site-packages/jax_unirep/__init__.py", line 1, in <module>
    from .evotuning import evotune, fit
  File "/mnt/lsf-nas-1/os-shared/anaconda3/envs/hannes_bio_embeddings/lib/python3.7/site-packages/jax_unirep/evotuning.py", line 8, in <module>
    from jax import grad, jit
  File "/mnt/lsf-nas-1/os-shared/anaconda3/envs/hannes_bio_embeddings/lib/python3.7/site-packages/jax/__init__.py", line 22, in <module>
    from .api import (
  File "/mnt/lsf-nas-1/os-shared/anaconda3/envs/hannes_bio_embeddings/lib/python3.7/site-packages/jax/api.py", line 37, in <module>
    from . import core
  File "/mnt/lsf-nas-1/os-shared/anaconda3/envs/hannes_bio_embeddings/lib/python3.7/site-packages/jax/core.py", line 31, in <module>
    from . import dtypes
  File "/mnt/lsf-nas-1/os-shared/anaconda3/envs/hannes_bio_embeddings/lib/python3.7/site-packages/jax/dtypes.py", line 31, in <module>
    from .lib import xla_client
  File "/mnt/lsf-nas-1/os-shared/anaconda3/envs/hannes_bio_embeddings/lib/python3.7/site-packages/jax/lib/__init__.py", line 60, in <module>
    from jaxlib import cusolver
ImportError: cannot import name 'cusolver' from 'jaxlib' (/mnt/lsf-nas-1/os-shared/anaconda3/envs/hannes_bio_embeddings/lib/python3.7/site-packages/jaxlib/__init__.py)

More info

[sacdallago]

Probably related to #161

[TheKursta]

Hi developers,

I'm currently encountering the same issue when importing UniRepEmbedder. Are there fixes or workarounds for this issue?

My environment:
python 3.7.11
bio_embeddings 0.2.2
jax 0.1.77
jax-unirep 1.0.3
jaxlib 0.1.75

[TheKursta]

Currently for me temporary fix has been to update jax packages in conda environment for bio_embeddings.
pip install jax==0.2.27
pip install jax-unirep==2.1.0
pip install jaxlib==0.1.76

[lam-ian]

Hi @TheKursta, I ran into same cusolver import error when attempting to use the UniRepEmbedder. I upgraded my jax packages to the same as yours but now ran into another error when trying to create a UniRepEmbedder object.

The error is as follows:
ImportError: cannot import name 'load_params_1900' from 'jax_unirep.utils' (/home/ec2-user/SageMaker/custom-miniconda/miniconda/envs/py38_bioembedding/lib/python3.8/site-packages/jax_unirep/utils.py)

So the jax_unirep.utils.py in jax-unirep v2.1.0 doesn't have a load_params_1900 function anymore. Just curious if you also ran into the similar issues after updating your jax libraries.

Environment:
Python 3.8
jax 0.2.28
jax-unirep 2.1.0
jaxlib 0.1.76

Traceback

---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
Input In [6], in <module>
----> 1 embedder = UniRepEmbedder()

File ~/SageMaker/custom-miniconda/miniconda/envs/py38_bioembedding/lib/python3.8/site-packages/bio_embeddings/embed/unirep_embedder.py:33, in UniRepEmbedder.__init__(self, device, **kwargs)
     32 def __init__(self, device: Union[None, str, torch.device] = None, **kwargs):
---> 33     from jax_unirep.utils import load_params_1900
     34     from jax_unirep.featurize import apply_fun
     36     self._params = load_params_1900()

ImportError: cannot import name 'load_params_1900' from 'jax_unirep.utils' (/home/ec2-user/SageMaker/custom-miniconda/miniconda/envs/py38_bioembedding/lib/python3.8/site-packages/jax_unirep/utils.py)