在Google运行时运行AlphaFold遇Config无jax_experimental_name_stack属性错误

错误原因 #

这个错误源于JAX与Haiku库的版本不兼容：旧版本Haiku代码尝试访问JAX的jax_experimental_name_stack配置项，但该属性在较新的JAX版本中已被移除或更名，导致属性查找失败。

解决方法 #

方法一：锁定兼容的依赖版本 #

在Colab的代码单元格中运行以下命令，重新安装指定版本的JAX、jaxlib和Haiku：

!pip install -q jax==0.4.13 jaxlib==0.4.13 haiku==0.0.9

安装完成后，点击Colab菜单栏的「Runtime」→「Restart runtime」重启运行时，再重新执行AlphaFold的预测代码即可。

方法二：临时修改Haiku源码 #

如果不想重新安装依赖，可直接修改报错的Haiku模块文件来绕过问题：
在Colab中运行以下命令，替换源码中失效的属性判断：

!sed -i 's/jax.config.jax_experimental_name_stack/jax.config.jax_numpy_rank_promotion is not None/g' /opt/conda/lib/python3.7/site-packages/haiku/_src/module.py

修改完成后无需重启运行时，直接重新运行AlphaFold的预测代码即可。

注：优先推荐方法一，版本锁定能从根源避免后续其他兼容性问题；方法二属于临时修复，仅适用于快速解决当前报错场景。

内容的提问来源于stack exchange，提问作者Advait R Aithal