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

[yuuu-d]

I installed jax-0.3.1 and jaxlib-0.3.0+cuda11.cudnn82 and tried to run Alphafold2. But I got an error, AttributeError: 'Config' object has no attribute 'jax_experimental_name_stack'. I‘ve no idea why this happened.
I spent a lot of time and I didn't find a similar problem.

[jakevdp]

jax_experimental_name_stack was added in #8395, and should be present in all jax releases with version 0.3.2 and newer. You report using jax version 0.3.1, which was released before this configuration was added.

[yuuu-d]

I was using the latest version in the past, but it reported an error that module 'jax.ops' has no attribute 'index_add'. After searching in the issue, I found that the version can be reduced to 0.3.1 to solve this problem. But jax_experimental_name_stack needs to be upgraded to 0.3.2 again. I upgraded the version to 0.3.2 and replaced index_add with x.at[idx].add(y) and finally solved the problem.
Thanks again for your help!

[hellojiangbo]

@ZhaoHanYuuu
How to i don't know replaced index_add with x.at[idx].add(y) .
Could you please help me ?

Thanks.

[jakevdp]

Hi - you should replace this:

jax.ops.index_add(x, idx, y)

with this:

x.at[idx].add(y)

Documentation here: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html