Can not run"predict the protein-ligand structure"

[Cathyliu99]

Hi Patrick, when I was trying Umol, I always meet ERROR with the following:
AttributeError Traceback (most recent call last)
in <cell line: 3>()
1 #@markdown #Predict the protein-ligand structure (a few minutes)
2 from net.model import config
----> 3 from predict_colab import predict
4 import numpy as np
5 MSA_FEATS='/content/'+ID+'/msa_features.pkl'

6 frames
/content/Umol/src/predict_colab.py in
4 import time
5 from typing import NamedTuple
----> 6 import haiku as hk
7 import jax
8 import jax.numpy as jnp

/usr/local/lib/python3.10/dist-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

/usr/local/lib/python3.10/dist-packages/haiku/experimental/init.py in
35 from haiku._src.dot import to_dot
36 from haiku._src.eval_shape import fast_eval_shape
---> 37 from haiku._src.layer_stack import layer_stack
38 from haiku._src.layer_stack import LayerStackTransparencyMapping
39 from haiku._src.lift import lift

/usr/local/lib/python3.10/dist-packages/haiku/_src/layer_stack.py in
21
22 from haiku._src import base
---> 23 from haiku._src import lift
24 from haiku._src import module
25 from haiku._src import transform

/usr/local/lib/python3.10/dist-packages/haiku/_src/lift.py in
22 from haiku._src import data_structures
23 from haiku._src import module
---> 24 from haiku._src import transform
25 from haiku._src.typing import LiftingModuleType
26

/usr/local/lib/python3.10/dist-packages/haiku/_src/transform.py in
340
341 COMPILED_FN_TYPES = (jax.lib.xla_extension.PjitFunction,
--> 342 jax.interpreters.xla.xe.PmapFunction) # pytype: disable=name-error
343
344

/usr/local/lib/python3.10/dist-packages/jax/_src/deprecations.py in getattr(name)
52 message, fn = deprecations[name]
53 if fn is None: # Is the deprecation accelerated?
---> 54 raise AttributeError(message)
55 warnings.warn(message, DeprecationWarning, stacklevel=2)
56 return fn

AttributeError: jax.interpreters.xla.xe was removed in JAX v0.4.36. Use jax.lib.xla_extension instead.

Could you please help me out?

[patrickbryant1]

Jax has been updated which causes errors, we will see if we can get a fix in colab soon.

Best,

Patrick