Jax Version and pocketfft

[BlankSama]

Facing issue while running alpha fold v2.2 and using jax==0.2.14 and jaxlib==0.3.10 and dm-haikuu==0.0.4

Traceback (most recent call last):
File "/home/datafiles/alphafold_data/alphafold/run_alphafold.py", line 33, in
from alphafold.model import data
File "/home/datafiles/alphafold_data/alphafold/alphafold/model/data.py", line 21, in
import haiku as hk
File "/home/conda1/anaconda3/envs/alphafold/lib/python3.8/site-packages/haiku/init.py", line 17, in
from haiku import data_structures
File "/home/conda1/anaconda3/envs/alphafold/lib/python3.8/site-packages/haiku/data_structures.py", line 17, in
from haiku._src.data_structures import to_immutable_dict
File "/home/conda1/anaconda3/envs/alphafold/lib/python3.8/site-packages/haiku/_src/data_structures.py", line 30, in
from haiku._src import utils
File "/home/conda1/anaconda3/envs/alphafold/lib/python3.8/site-packages/haiku/_src/utils.py", line 24, in
import jax
File "/home/conda1/anaconda3/envs/alphafold/lib/python3.8/site-packages/jax/init.py", line 108, in
from .experimental.maps import soft_pmap
File "/home/conda1/anaconda3/envs/alphafold/lib/python3.8/site-packages/jax/experimental/maps.py", line 25, in
from .. import numpy as jnp
File "/home/conda1/anaconda3/envs/alphafold/lib/python3.8/site-packages/jax/numpy/init.py", line 16, in
from . import fft
File "/home/conda1/anaconda3/envs/alphafold/lib/python3.8/site-packages/jax/numpy/fft.py", line 17, in
from jax._src.numpy.fft import (
File "/home/conda1/anaconda3/envs/alphafold/lib/python3.8/site-packages/jax/_src/numpy/fft.py", line 19, in
from jax import lax
File "/home/conda1/anaconda3/envs/alphafold/lib/python3.8/site-packages/jax/lax/init.py", line 330, in
from jax._src.lax.fft import (
File "/home/conda1/anaconda3/envs/alphafold/lib/python3.8/site-packages/jax/_src/lax/fft.py", line 144, in
xla.backend_specific_translations['cpu'][fft_p] = pocketfft.pocketfft
AttributeError: module 'jaxlib.pocketfft' has no attribute 'pocketfft'

Please help

[SimonCK666]

Looking for the solution, i got the same problem

[baba-hashimoto]

google-deepmind/alphafold#513

[sokrypton]

You get this error when there is a mismatch between jax and jaxlib. You may need to downgrade your jaxlib.

[karlaarz]

Hello, is there any solution for this? I tried this solution google-deepmind/alphafold#513 but still get same error.

[BlankSama]

@karlaarz yeah i found the solution for it jaxlib == 0.1.69+cuda111 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html and jax version= 0.2.14