New Jax version broke Alphafold on Colab?
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)

[<ipython-input-5-420c51c385fa>](https://localhost:8080/#) in <module>
     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](https://localhost:8080/#) 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' only work with multiple sequences,if I run with only  one sequence, "AttributeError: module 'jax' has no attribute 'tree_multimap'" comes out 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? 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). I just successfully used a different Colab notebook: 
https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/AlphaFold2.ipynb 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. HI! Thanks for raising this issue. We're currently testing a fix and should hopefully be able to push an update shortly. Thanks! 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! Hi thanks for this! This has now been fixed in https://github.com/deepmind/alphafold/releases/tag/v2.2.4. I'm getting another error now, at the database search stage:

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

![image](https://user-images.githubusercontent.com/18224663/191625792-67ed3d11-863e-411b-9635-c5b7e9546244.png)
 Did you execute cell 3 before cell 4? This is where model_type_to_use is defined. Thanks! Doh thanks, yes I was a a bit too hurried.
Was getting deprecation warnings for a while, now appears not to work at all.

![image](https://user-images.githubusercontent.com/18224663/188844908-3b63f131-babb-4efb-82de-08aeb573a36a.png)
