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

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'