AttributeError: module 'jax' has no attribute 'tree_multimap'
Hi thanks for raising this! This duplicates https://github.com/deepmind/alphafold/issues/589 so will close this issue Thank you so much for getting back.
This is related to #[456 ](https://github.com/deepmind/alphafold/issues/456). 

"5. Run AlphaFold and download prediction" (in Colab notebook [example](https://colab.research.google.com/github/deepmind/alphafold/blob/main/notebooks/AlphaFold.ipynb)), started throwing an error recently and now expects 'jax.tree_multimap' to be changed to 'jax.tree_map' (see [stackflow](https://stackoverflow.com/questions/73693738/module-jax-has-no-attribute-tree-multimap-in-alphafold2-colab))

[tfgg](https://github.com/tfgg) noted this back in May but maybe we should update this to accommodate the latest jax module in Colab. (jax 0.3.17 jaxlib 0.3.15+cuda11.cudnn805)