Alpha Fold crash
Same error here but unsure how to solve it This is an upgrade from google, the newest version of colab by default takes up jax 0.3.25 (https://colab.research.google.com/notebooks/relnotes.ipynb#scrollTo=iOWrh41ocOca) which has no support for 'jax_experimental_name_stack', I have solved this issue by downgrading colab jax to 0.3.17, you can do it by adding this command: !pip install --upgrade "jax[cuda]==0.3.17" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html and then restarting the runtime > This is an upgrade from google, the newest version of colab by default takes up jax 0.3.25 (https://colab.research.google.com/notebooks/relnotes.ipynb#scrollTo=iOWrh41ocOca) which has no support for 'jax_experimental_name_stack', I have solved this issue by downgrading colab jax to 0.3.17, you can do it by adding this command: !pip install --upgrade "jax[cuda]==0.3.17" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html and then restarting the runtime

Hey, I have been struggeling to apply your solution. Could you elaborate further on when and where to insert this command? Thanks for raising this. Closing as this is a duplicate of https://github.com/deepmind/alphafold/issues/635.
After running AF Colab multimer for weeks, it started to get unstable this week. Recently it has crashed multiple times. Each when starting task 5 (Run AlphaFold and download prediction) the system crashes within no time although I am Prolab+ user. 
The following or similar error message occurs, which is similar to what other users feedback:

```
/opt/conda/lib/python3.7/site-packages/haiku/_src/data_structures.py:144: FutureWarning: jax.tree_flatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_flatten instead.
  leaves, treedef = jax.tree_flatten(tree)
/opt/conda/lib/python3.7/site-packages/haiku/_src/data_structures.py:145: 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(treedef, leaves)
---------------------------------------------------------------------------
UnfilteredStackTrace                      Traceback (most recent call last)
[<ipython-input-5-ca4ee2dc266d>](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 

________________________________________
23 frames
________________________________________
UnfilteredStackTrace: AttributeError: 'Config' object has no attribute 'jax_experimental_name_stack'

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/module.py](https://localhost:8080/) in wrapped(self, *args, **kwargs)
    406       f = functools.partial(unbound_method, self)
    407       f = functools.partial(run_interceptors, f, method_name, self)
--> 408       if jax.config.jax_experimental_name_stack and module_name:
    409         local_module_name = module_name.split("/")[-1]
    410         f = jax.named_call(f, name=local_module_name)

AttributeError: 'Config' object has no attribute 'jax_experimental_name_stack'
```