AttributeError: 'Config' object has no attribute 'jax_experimental_name_stack'
Hi! I'm getting the same error. It seems like Colab version of AlphaFold uses JAX v0.3.25. `jax.config.jax_experimental_name_stack` was removed in this version. If you manage to use older version of JAX with this Colab notebook, I guess it would work. Hi! I'm also receiving this error. The AlphaFold `requirements.txt` file which is installed via pip3 in cell 2 specifies JAX 0.3.17, so I'm unsure why the version changes (running `print(jax.__version__)` at the start of cell 5 prints `0.3.25`) +1 using a monomer.

Full error message 
![image](https://user-images.githubusercontent.com/18224663/203679765-bafc3440-28fb-465d-a96c-3895f428964c.png)
 Similar problem and it only started this week.  Hello. Same issue since november 18. I believe I have the same issue, first noticed 2-3 days ago with DeepMind's AlphaFold Colab (now running AF2.2.4), when I try to run it in multimer form, without exceeding sequence size limits. Things proceed well through construction of MSAs, but as it starts to run Model_1_multimer, it crashes with the following error message (below). Thx in advance for your kind help!

/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-b59661ebfbb4> 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 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 Many thanks for the information. I thought s.th. must have happened.
Best wishes,
 Heinz



Dr Heinz Peter Nasheuer

Personal Professor

School of Biological and Chemical Sciences

Biochemistry, Room 110

Arts & Science Building, Main Concourse

University of Galway

Distillery Road

Galway, H91 TK33, Ireland


Phone: ++353-91-49 2430

________________________________
From: anupam-tiwari ***@***.***>
Sent: 26 November 2022 14:58
To: deepmind/alphafold ***@***.***>
Cc: Nasheuer, Heinz-Peter ***@***.***>; Comment ***@***.***>
Subject: Re: [deepmind/alphafold] AttributeError: 'Config' object has no attribute 'jax_experimental_name_stack' (Issue #635)

EXTERNAL EMAIL: This email originated outside the University of Galway. Do not open attachments or click on links unless you believe the content is safe.
RÍOMHPHOST SEACHTRACH: Níor tháinig an ríomhphost seo ó Ollscoil na Gaillimhe. Ná hoscail ceangaltáin agus ná cliceáil ar naisc mura gcreideann tú go bhfuil an t-ábhar sábháilte.


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

—
Reply to this email directly, view it on GitHub<https://github.com/deepmind/alphafold/issues/635#issuecomment-1328060690>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/A4LJE65LQUP4OSICZQPNDUDWKIQQJANCNFSM6AAAAAASIFO6AQ>.
You are receiving this because you commented.Message ID: ***@***.***>
 > 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? I can't get it to work either. jax version appears to already 0.3.17 while jaxlib is 0.3.25. But I can't seem to install 0.3.17.

![image](https://user-images.githubusercontent.com/18224663/204108000-881279a7-157c-4e4a-aead-b5041afb0e59.png)
![image](https://user-images.githubusercontent.com/18224663/204108036-1853c7df-6ddb-4044-9099-307f1f6bdc17.png)
 The earliest version I was able to downgrade jaxlib to is 0.3.20, however this still leads to the original error. > > 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?

Same I have been having issue where to run the command that was provided. If someone could help with this, that would be great! https://colab.research.google.com/drive/1J1G8OC6LZwT6Tqdz4uLRmnVzj7WWgUg1?usp=sharing 
here you go guys, I made a modified notebook, the first step involves jax and jaxlib step, let me know how it goes!  Yes, it seems to work. Thank you!

I don't quite understand why you downgrade *before* downloading the packages, but it works.  I tried to include this in the downloading package step, but to use the downgraded version we have to restart the runtime which in turn messes up the step-1 and step-2 code, hence I made it a separate step Thanks for doing this!

I tried but now I got a different error message. 


NameError                            Traceback (most recent call last)
<ipython-input-2-15b941f78883><https://localhost:8080/> in <module>
     18
     19 # --- Run the model ---
---> 20 if model_type_to_use == notebook_utils.ModelType.MONOMER:
     21   model_names = config.MODEL_PRESETS['monomer'] + ('model_2_ptm',)
     22 elif model_type_to_use == notebook_utils.ModelType.MULTIMER:

NameError: name 'model_type_to_use' is not defined


Just to confirm I used the multimer model with 2 sequences and they were run fine through the MSA step. > Hello. Same issue since november 18.

Agree Nov 18th in the morning everything was working fine and in the new week it started. Thanks for this! Closing as this is a duplicate of https://github.com/deepmind/alphafold/issues/634 I appreciate all the comments! The issue has been solved in the new version! I just run two chain prediction successfully. Many thanks!
Hello,

I am predicting protein multimers using AlphaFold2 however it always failed and showed something like the following:

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'


It shows "Attribute error" but I have no idea about coding to solve this issue.
I would really appreciate if someone knows what is going on.

Tons of thanks!
Kaka