New Jax version broke Alphafold on Colab?

[danjaco]

Was getting deprecation warnings for a while, now appears not to work at all.

[Mila0007]

I am getting the same error message. Just in case it might be related to this problem, I was able to obtain protein predictions before the relax_use_gpu option was included in the notebook.

Error message:

/opt/conda/lib/python3.7/site-packages/haiku/_src/data_structures.py:195: FutureWarning: jax.tree_flatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_flatten instead. leaves, structure = jax.tree_flatten(mapping) /opt/conda/lib/python3.7/site-packages/haiku/_src/data_structures.py:203: FutureWarning: jax.tree_unflatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_unflatten instead. self._mapping = jax.tree_unflatten(self._structure, self._leaves) /opt/conda/lib/python3.7/site-packages/alphafold/model/mapping.py:50: FutureWarning: jax.tree_flatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_flatten instead. values_tree_def = jax.tree_flatten(values)[1] /opt/conda/lib/python3.7/site-packages/alphafold/model/mapping.py:54: FutureWarning: jax.tree_unflatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_unflatten instead. return jax.tree_unflatten(values_tree_def, flat_axes) /opt/conda/lib/python3.7/site-packages/alphafold/model/mapping.py:129: FutureWarning: jax.tree_flatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_flatten instead. flat_sizes = jax.tree_flatten(in_sizes)[0] /opt/conda/lib/python3.7/site-packages/haiku/_src/stateful.py:457: FutureWarning: jax.tree_leaves is deprecated, and will be removed in a future release. Use jax.tree_util.tree_leaves instead. length = jax.tree_leaves(xs)[0].shape[0]

---

UnfilteredStackTrace Traceback (most recent call last)

in 44 processed_feature_dict = model_runner.process_features(np_example, random_seed=0) ---> 45 prediction = model_runner.predict(processed_feature_dict, random_seed=random.randrange(sys.maxsize)) 46

87 frames

UnfilteredStackTrace: AttributeError: module 'jax' has no attribute 'tree_multimap'

The stack trace below excludes JAX-internal frames. The preceding is the original exception that occurred, unmodified.

---

The above exception was the direct cause of the following exception:

AttributeError Traceback (most recent call last)

/opt/conda/lib/python3.7/site-packages/haiku/_src/stateful.py in difference(before, after) 310 params_before, params_after = box_and_fill_missing(before.params, 311 after.params) --> 312 params_after = jax.tree_multimap(functools.partial(if_changed, is_new_param), 313 params_before, params_after) 314

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

[WS-FL]

only work with multiple sequences,if I run with only one sequence, "AttributeError: module 'jax' has no attribute 'tree_multimap'" comes out

[ruth-hanna]

I am having this issue as well when trying to run single sequences in the colab notebook. Has anyone found a workaround that allows notebooks to run to completion?

[Mila0007]

It doesn't look as if someone has suggested a workaround. I tried again today using a single sequence and I am still getting the same error message reported a week ago. Hopefully someone will look into this issue as soon as possible, as it has completely broken the Colab notebook (at least when running single sequences).

[ruth-hanna]

I just successfully used a different Colab notebook: https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/AlphaFold2.ipynb

[Mila0007]

Thank you very much for letting us know @ruth-hanna! Hopefully this issue will be solved soon and we will be able to use this notebook as well.

[Htomlinson14]

HI! Thanks for raising this issue. We're currently testing a fix and should hopefully be able to push an update shortly. Thanks!

[janinelux]

Is there any update on the 'jax' issue for https://colab.research.google.com/github/deepmind/alphafold/blob/main/notebooks/AlphaFold.ipynb ? I also have this issue for predictions from one protein sequence. Thank you for trying to fix it!

[Htomlinson14]

Hi thanks for this! This has now been fixed in https://github.com/deepmind/alphafold/releases/tag/v2.2.4.

[danjaco]

I'm getting another error now, at the database search stage:

"NameError: name 'model_type_to_use' is not defined" (line 80)

[Htomlinson14]

Did you execute cell 3 before cell 4? This is where model_type_to_use is defined. Thanks!

[danjaco]

Doh thanks, yes I was a a bit too hurried.