ModuleNotFoundError at "Run Prediction Step"

[william15961]

Keep getting this error when attempting to run prediction.

ModuleNotFoundError Traceback (most recent call last) in () 5 from colabfold.download import download_alphafold_params, default_data_dir 6 from colabfold.utils import setup_logging ----> 7 from colabfold.batch import get_queries, run, set_model_type 8 K80_chk = get_ipython().getoutput('nvidia-smi | grep "Tesla K80" | wc -l') 9 if "1" in K80_chk:

9 frames /usr/local/lib/python3.7/dist-packages/jax/_src/lib/mlir/dialects/init.py in () 17 import jaxlib.mlir.dialects.chlo as chlo 18 import jaxlib.mlir.dialects.mhlo as mhlo ---> 19 import jaxlib.mlir.dialects.std as std

ModuleNotFoundError: No module named 'jaxlib.mlir.dialects.std'

[milot-mirdita]

Please restart (and "factory reset") your ColabFold colab. I added a workaround that installs an older jax version.

[ExaptedAI]

After 'restart' and 'factory reset' I still get the following error when running AlphaFold2 Batch.

ModuleNotFoundError Traceback (most recent call last) in () 3 import sys 4 ----> 5 from colabfold.batch import get_queries, run 6 from colabfold.download import default_data_dir 7 from colabfold.utils import setup_logging

9 frames /usr/local/lib/python3.7/dist-packages/jax/_src/lib/mlir/dialects/init.py in () 17 import jaxlib.mlir.dialects.chlo as chlo 18 import jaxlib.mlir.dialects.mhlo as mhlo ---> 19 import jaxlib.mlir.dialects.std as std

ModuleNotFoundError: No module named 'jaxlib.mlir.dialects.std'

[milot-mirdita]

The batch notebook should now be fixed too.

[ExaptedAI]

Perfect! Thank you very much Milot. It works fine now. I appreciate the fast turnaround on the fix.

Looking in links: https://storage.googleapis.com/jax-releases/jax_releases.html Requirement already satisfied: jax[cuda]<0.3.0 in /usr/local/lib/python3.7/dist-packages (0.2.28) Requirement already satisfied: scipy>=1.2.1 in /usr/local/lib/python3.7/dist-packages (from jax[cuda]<0.3.0) (1.4.1) Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from jax[cuda]<0.3.0) (3.10.0.2) Requirement already satisfied: opt-einsum in /usr/local/lib/python3.7/dist-packages (from jax[cuda]<0.3.0) (3.3.0) Requirement already satisfied: numpy>=1.19 in /usr/local/lib/python3.7/dist-packages (from jax[cuda]<0.3.0) (1.21.5) Requirement already satisfied: absl-py in /usr/local/lib/python3.7/dist-packages (from jax[cuda]<0.3.0) (0.13.0) Collecting jaxlib==0.1.76+cuda11.cudnn82 Downloading https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.1.76%2Bcuda11.cudnn82-cp37-none-manylinux2010_x86_64.whl (153.1 MB) Requirement already satisfied: flatbuffers<3.0,>=1.12 in /usr/local/lib/python3.7/dist-packages (from jaxlib==0.1.76+cuda11.cudnn82->jax[cuda]<0.3.0) (2.0) Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from absl-py->jax[cuda]<0.3.0) (1.15.0) Installing collected packages: jaxlib Attempting uninstall: jaxlib Found existing installation: jaxlib 0.3.2+cuda11.cudnn805 Uninstalling jaxlib-0.3.2+cuda11.cudnn805: Successfully uninstalled jaxlib-0.3.2+cuda11.cudnn805 Successfully installed jaxlib-0.1.76+cuda11.cudnn82 Downloading alphafold2 weights to /root/.cache/colabfold: 100%|██████████| 3.82G/3.82G [00:24<00:00, 165MB/s] Downloading alphafold2 weights to /root/.cache/colabfold: 100%|██████████| 3.47G/3.47G [00:30<00:00, 124MB/s]