AttributeError: module 'jax' has no attribute 'linear_util'

[laurahum]

I'm running alphafold colab in Firefox browser and when running the section '4. Search against genetic databases' I get the following error:

AttributeError Traceback (most recent call last)

in () 25 import py3Dmol 26 ---> 27 from alphafold.model import model 28 from alphafold.model import config 29 from alphafold.model import data 7 frames /opt/conda/lib/python3.10/site-packages/alphafold/model/model.py in 19 from alphafold.common import confidence 20 from alphafold.model import features ---> 21 from alphafold.model import modules 22 from alphafold.model import modules_multimer 23 import haiku as hk /opt/conda/lib/python3.10/site-packages/alphafold/model/modules.py in 19 import functools 20 from alphafold.common import residue_constants ---> 21 from alphafold.model import all_atom 22 from alphafold.model import common_modules 23 from alphafold.model import folding /opt/conda/lib/python3.10/site-packages/alphafold/model/all_atom.py in 38 39 from alphafold.model import r3 ---> 40 from alphafold.model import utils 41 import jax 42 import jax.numpy as jnp /opt/conda/lib/python3.10/site-packages/alphafold/model/utils.py in 21 from typing import Mapping 22 ---> 23 import haiku as hk 24 import jax 25 import jax.numpy as jnp /opt/conda/lib/python3.10/site-packages/haiku/__init__.py in 18 from haiku import config 19 from haiku import data_structures ---> 20 from haiku import experimental 21 from haiku import initializers 22 from haiku import mixed_precision /opt/conda/lib/python3.10/site-packages/haiku/experimental/__init__.py in 32 from haiku._src.config import module_auto_repr 33 from haiku._src.config import rng_reserve_size ---> 34 from haiku._src.dot import abstract_to_dot 35 from haiku._src.dot import to_dot 36 from haiku._src.eval_shape import fast_eval_shape /opt/conda/lib/python3.10/site-packages/haiku/_src/dot.py in 161 162 --> 163 @jax.linear_util.transformation 164 def _interpret_subtrace(main, *in_vals): 165 trace = DotTrace(main, jax.core.cur_sublevel()) /usr/local/lib/python3.10/dist-packages/jax/_src/deprecations.py in getattr(name) 52 warnings.warn(message, DeprecationWarning, stacklevel=2) 53 return fn ---> 54 raise AttributeError(f"module {module!r} has no attribute {name!r}") 55 56 return getattr AttributeError: module 'jax' has no attribute 'linear_util'

[ScottStro]

Hi, Have you found a workaround for this issue? Thanks

[KesharwaniArpita]

I am facing the same issue in my research.

[xianyu-123]

An error free installation package https://github.com/google-deepmind/alphafold/issues/933

[LincolnG4]

On Alphafold's folder go to docker/Dockerfile and downgrade jax (line 70) and jaxlib (line 71) to version 0.4.23

jax==0.4.23 \
jaxlib==0.4.23+cuda12.cudnn89 \

after that, rebuild the docker image (docker build -f docker/Dockerfile -t alphafold . ) it may help