AlphaFold Colab crashing at cell 5
This is the full Stack Trace:
```
/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
[/opt/conda/lib/python3.7/site-packages/alphafold/model/model.py](https://localhost:8080/#) in predict(self, feat, random_seed)
    166                  tree.map_structure(lambda x: x.shape, feat))
--> 167     result = self.apply(self.params, jax.random.PRNGKey(random_seed), feat)
    168 

[/usr/local/lib/python3.7/dist-packages/jax/_src/traceback_util.py](https://localhost:8080/#) in reraise_with_filtered_traceback(*args, **kwargs)
    161     try:
--> 162       return fun(*args, **kwargs)
    163     except Exception as e:

[/usr/local/lib/python3.7/dist-packages/jax/_src/api.py](https://localhost:8080/#) in cache_miss(*args, **kwargs)
    621         jax.config.jax_debug_nans or jax.config.jax_debug_infs):
--> 622       execute = dispatch._xla_call_impl_lazy(fun_, *tracers, **params)
    623       out_flat = call_bind_continuation(execute(*args_flat))

[/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py](https://localhost:8080/#) in _xla_call_impl_lazy(***failed resolving arguments***)
    236   return xla_callable(fun, device, backend, name, donated_invars, keep_unused,
--> 237                       *arg_specs)
    238 

[/usr/local/lib/python3.7/dist-packages/jax/linear_util.py](https://localhost:8080/#) in memoized_fun(fun, *args)
    302     else:
--> 303       ans = call(fun, *args)
    304       cache[key] = (ans, fun.stores)

[/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py](https://localhost:8080/#) in _xla_callable_uncached(fun, device, backend, name, donated_invars, keep_unused, *arg_specs)
    359     return lower_xla_callable(fun, device, backend, name, donated_invars, False,
--> 360                               keep_unused, *arg_specs).compile().unsafe_call
    361 

[/usr/local/lib/python3.7/dist-packages/jax/_src/profiler.py](https://localhost:8080/#) in wrapper(*args, **kwargs)
    313     with TraceAnnotation(name, **decorator_kwargs):
--> 314       return func(*args, **kwargs)
    315     return wrapper

[/usr/local/lib/python3.7/dist-packages/jax/_src/dispatch.py](https://localhost:8080/#) in lower_xla_callable(fun, device, backend, name, donated_invars, always_lower, keep_unused, *arg_specs)
    445     jaxpr, out_type, consts = pe.trace_to_jaxpr_final2(
--> 446         fun, pe.debug_info_final(fun, "jit"))
    447   out_avals, kept_outputs = util.unzip2(out_type)

[/usr/local/lib/python3.7/dist-packages/jax/_src/profiler.py](https://localhost:8080/#) in wrapper(*args, **kwargs)
    313     with TraceAnnotation(name, **decorator_kwargs):
--> 314       return func(*args, **kwargs)
    315     return wrapper

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/partial_eval.py](https://localhost:8080/#) in trace_to_jaxpr_final2(fun, debug_info)
   2076     with core.new_sublevel():
-> 2077       jaxpr, out_type, consts = trace_to_subjaxpr_dynamic2(fun, main, debug_info)
   2078     del fun, main

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/partial_eval.py](https://localhost:8080/#) in trace_to_subjaxpr_dynamic2(fun, main, debug_info)
   2026     in_tracers_ = [t for t, keep in zip(in_tracers, keep_inputs) if keep]
-> 2027     ans = fun.call_wrapped(*in_tracers_)
   2028     out_tracers = map(trace.full_raise, ans)

[/usr/local/lib/python3.7/dist-packages/jax/linear_util.py](https://localhost:8080/#) in call_wrapped(self, *args, **kwargs)
    166     try:
--> 167       ans = self.f(*args, **dict(self.params, **kwargs))
    168     except:

[/opt/conda/lib/python3.7/site-packages/haiku/_src/transform.py](https://localhost:8080/#) in apply_fn(params, *args, **kwargs)
    126 
--> 127     out, state = f.apply(params, {}, *args, **kwargs)
    128     if state:

[/opt/conda/lib/python3.7/site-packages/haiku/_src/transform.py](https://localhost:8080/#) in apply_fn(params, state, rng, *args, **kwargs)
    353       try:
--> 354         out = f(*args, **kwargs)
    355       except jax.errors.UnexpectedTracerError as e:

[/opt/conda/lib/python3.7/site-packages/alphafold/model/model.py](https://localhost:8080/#) in _forward_fn(batch)
     81       def _forward_fn(batch):
---> 82         model = modules.AlphaFold(self.config.model)
     83         return model(

[/opt/conda/lib/python3.7/site-packages/haiku/_src/module.py](https://localhost:8080/#) in __call__(cls, *args, **kwargs)
    120     init = wrap_method("__init__", cls.__init__)
--> 121     init(module, *args, **kwargs)
    122 

[/opt/conda/lib/python3.7/site-packages/haiku/_src/module.py](https://localhost:8080/#) in wrapped(self, *args, **kwargs)
    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]

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)
[<ipython-input-5-ca4ee2dc266d>](https://localhost:8080/#) in <module>
     43     model_runner = model.RunModel(cfg, params)
     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 
     47     mean_plddt = prediction['plddt'].mean()

[/opt/conda/lib/python3.7/site-packages/alphafold/model/model.py](https://localhost:8080/#) in predict(self, feat, random_seed)
    165     logging.info('Running predict with shape(feat) = %s',
    166                  tree.map_structure(lambda x: x.shape, feat))
--> 167     result = self.apply(self.params, jax.random.PRNGKey(random_seed), feat)
    168 
    169     # This block is to ensure benchmark timings are accurate. Some blocking is

[/opt/conda/lib/python3.7/site-packages/haiku/_src/transform.py](https://localhost:8080/#) in apply_fn(params, *args, **kwargs)
    125           "name (e.g. `f.apply(.., state=my_state)`)")
    126 
--> 127     out, state = f.apply(params, {}, *args, **kwargs)
    128     if state:
    129       raise ValueError("If your transformed function uses `hk.{get,set}_state` "

[/opt/conda/lib/python3.7/site-packages/haiku/_src/transform.py](https://localhost:8080/#) in apply_fn(params, state, rng, *args, **kwargs)
    352     with base.new_context(params=params, state=state, rng=rng) as ctx:
    353       try:
--> 354         out = f(*args, **kwargs)
    355       except jax.errors.UnexpectedTracerError as e:
    356         raise jax.errors.UnexpectedTracerError(unexpected_tracer_hint) from e

[/opt/conda/lib/python3.7/site-packages/alphafold/model/model.py](https://localhost:8080/#) in _forward_fn(batch)
     80     else:
     81       def _forward_fn(batch):
---> 82         model = modules.AlphaFold(self.config.model)
     83         return model(
     84             batch,

[/opt/conda/lib/python3.7/site-packages/haiku/_src/module.py](https://localhost:8080/#) in __call__(cls, *args, **kwargs)
    119     # Now attempt to initialize the object.
    120     init = wrap_method("__init__", cls.__init__)
--> 121     init(module, *args, **kwargs)
    122 
    123     if (config.get_config().module_auto_repr and

[/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'
``` 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 Thanks for raising this. This has been fixed in the latest commit.
Today I've tried to run Colab version of AlphaFold v2.2.4, but it crashed every time while executing cell number 5. I'm getting these errors:

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

**'Config' object has no attribute 'jax_experimental_name_stack'**