ModuleNotFoundError: No module named 'jax.extend' related to #209, #210

[ZubuNoShoshinsha]

Hello,

My question is related to #209 and #210 My environment is...

Wsl2 OS: Ubuntu 22.04.4 GCC: 11.4.0 CUDA: 12.1 GPU: RTX 4090 LocalColabFold Ver. 1.5.5

As instructed in #209 , I checked if GPU was recognized and it was not. So, I dongraded jax and jaxlib to jax 0.4.7 jaxlib0.4.7+cuda11.cudnn82 as instructed in #209 .

And then I checked again using $ /path/to/your/localcolabfold/colabfold-conda/bin/python3.10

import jax print(jax.local_devices()[0].platform)

and "gpu" was returned.

Then, I run the localcolabfold. But, this error message popped up and stopped like below

2024-04-01 15:14:35,452 Running colabfold 1.5.5 (61df3b853140ca79dbdf64349824beb14364ebfd) 2024-04-01 15:14:36,006 Running on GPU Traceback (most recent call last): File "/mnt/d/Alphafold/localcolabfold/colabfold-conda/bin/colabfold_batch", line 8, in sys.exit(main()) File "/mnt/d/AlphaFold/localcolabfold/colabfold-conda/lib/python3.10/site-packages/colabfold/batch.py", line 2037, in main run( File "/mnt/d/AlphaFold/localcolabfold/colabfold-conda/lib/python3.10/site-packages/colabfold/batch.py", line 1292, in run from colabfold.alphafold.models import load_models_and_params File "/mnt/d/AlphaFold/localcolabfold/colabfold-conda/lib/python3.10/site-packages/colabfold/alphafold/models.py", line 4, in import haiku File "/mnt/d/AlphaFold/localcolabfold/colabfold-conda/lib/python3.10/site-packages/haiku/init.py", line 20, in from haiku import experimental File "/mnt/d/AlphaFold/localcolabfold/colabfold-conda/lib/python3.10/site-packages/haiku/experimental/init.py", line 34, in from haiku._src.dot import abstract_to_dot File "/mnt/d/AlphaFold/localcolabfold/colabfold-conda/lib/python3.10/site-packages/haiku/_src/dot.py", line 29, in from jax.extend import linear_util as lu ModuleNotFoundError: No module named 'jax.extend'

It would be helpful if there would be any instruction for solving this issue.

[YoshitakaMo]

I suspect that the issue lies in the version of the dm-haiku module being 0.0.11 or later. In my environment:

$ localcolabfold/colabfold-conda/bin/python3.10 -m pip list

jax                          0.4.23
jaxlib                       0.4.23+cuda11.cudnn86
chex                         0.1.85
dm-haiku                     0.0.10

If CUDA 12.1 is installed, these versions should be fine.

Please set your dm-haiku to version 0.0.10. Otherwise, you may encounter the error ModuleNotFoundError: No module named 'jax.extend'

[ZubuNoShoshinsha]

Thank you for your suggestion. (I just noticed your response) Actually my dm-haiku was 0.0.12, so I down graded to 0.0.10 as you suggested.

And I ran localcolabfold 1.5.5. So, my environment is now jax 0.4.7 jaxlib 0.4.7+cuda11.cudnn82 chex 0.1.82 dm-haiku 0.0.10

No more " ModuleNotFoundError: No module named 'jax.extend' ", but now new message showed up and the program stopped.

" Could not predict ProteinA. Not Enough GPU memory? FAILED_PRECONDITION: DNN library initialization failed. Look at the errors above for more details. "

Is there any more suggestion to solve this??

[ZubuNoShoshinsha]

my nvidia-smi result

[YoshitakaMo]

If you are using WSL2, did you turn on the settings shown in https://github.com/YoshitakaMo/localcolabfold?tab=readme-ov-file#for-wsl2-in-windows ? Unfortunately, I can't figure out the cause because I don't have a WSL2 environment.

[ZubuNoShoshinsha]

Yes, I did. I restart wsl2 and tried another shot, but it didn't work well.

Thank you though.

[ZubuNoShoshinsha]

I wonder... when I downgraded dm-haiku, the message said " colabfold 1.5.5 requires dm-haiku<0.013, >=0.0.12, but you have dm-haiku 0.0.10 which is incompatible. " Is it fine to run colabfold appropriately?

[ZubuNoShoshinsha]

Finally,

I might have found the solution.

I downgraded " nvidia-cudnn-cu11 " by doing this command from 9.0.0.312 to 8.5.0.96 .

pip install --upgrade nvidia-cudnn-cu11==8.5.0.96

I ran the localcolabfold and it processed very smoothly on GPU.

I was astonished.

Thank you.

[Vinaysukhesh98]

jax                          0.4.23
jaxlib                       0.4.23+cuda11.cudnn86
chex                         0.1.85
dm-haiku                     0.0.10

Requirement already satisfied: torch==1.13.1 in /usr/local/lib/python3.10/dist-packages (1.13.1) Requirement already satisfied: transformers==4.24.0 in /usr/local/lib/python3.10/dist-packages (4.24.0) Collecting diffusers==0.3.0 Using cached diffusers-0.3.0-py3-none-any.whl (153 kB) Collecting jax==0.4.23 Downloading jax-0.4.23-py3-none-any.whl (1.7 MB) ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.7/1.7 MB 6.3 MB/s eta 0:00:00 ERROR: Could not find a version that satisfies the requirement jaxlib==0.4.23+cuda11.cudnn86 (from versions: 0.4.6, 0.4.7, 0.4.9, 0.4.10, 0.4.11, 0.4.12, 0.4.13, 0.4.14, 0.4.16, 0.4.17, 0.4.18, 0.4.19, 0.4.20, 0.4.21, 0.4.22, 0.4.23, 0.4.24, 0.4.25, 0.4.26, 0.4.27, 0.4.28) ERROR: No matching distribution found for jaxlib==0.4.23+cuda11.cudnn86

[YoshitakaMo]

I updated the installer and updater script for Linux two days ago as Jax 0.4.23 no longer seems suitable for cuda 12 and cudnn 9. Please update your cuda to 12.4, cudnn to 9, and use the latest updater script.