Alphafold2 Colab运行报错：module 'jax' has no attribute 'xla'求助

使用AlphaFold2时反复触发以下错误：

AttributeError                            Traceback (most recent call last)
<ipython-input-5-74c943a59072> in <cell line: 55>()
     58     runner = cf_af.prep_model_runner(opt, old_runner=runner)
     59   else:
---&gt; 60     runner = cf_af.prep_model_runner(opt)
     61 else:
     62   runner = None

________________________________________
7 frames
________________________________________
/usr/local/lib/python3.10/dist-packages/jax/_src/deprecations.py in getattr(name)
     51       warnings.warn(message, DeprecationWarning, stacklevel=2)
     52       return fn
---&gt; 53     raise AttributeError(f"module {module!r} has no attribute {name!r}")
     54 
     55   return getattr

AttributeError: module 'jax' has no attribute 'xla'

首次运行成功生成蛋白模型后开始报错，当前环境：Python 3.11.4、jax 0.4.14、haiku 0.0.5，曾短暂恢复并生成3个结构，但之后再次报错无法解决。尝试执行pip install "jax<=0.3.16" "jaxlib<=0.3.16"降级jax版本时，出现以下错误：

Using cached jax-0.3.16-py3-none-any.whl
ERROR: Could not find a version that satisfies the requirement jaxlib<=0.3.16 (from versions: 0.4.13, 0.4.14)
ERROR: No matching distribution found for jaxlib<=0.3.16

Python 3.11不兼容jax 0.3.x系列版本，jaxlib 0.3.x未提供适配Python 3.11的安装包，因此降级操作失败。而jax 0.4.x中jax.xla模块已被移除，AlphaFold2旧版本代码仍在调用该废弃模块，导致AttributeError。

方案1：使用AlphaFold2官方Docker镜像（推荐，无需手动处理依赖） #

Docker镜像已预配置所有兼容的依赖版本，彻底避免版本冲突：

1. 安装Docker（根据你的操作系统完成安装）
2. 拉取官方镜像：

   docker pull deepmind/alphafold
   

3. 按照官方指引，将你的蛋白序列文件挂载到容器中执行模型预测。

方案2：修复当前Python 3.11环境的依赖与代码问题 #

1. 卸载当前jax/jaxlib：

   pip uninstall jax jaxlib -y
   

2. 安装适配Python 3.11的jax版本（GPU用户需指定CUDA版本，CPU用户去掉+cuda11.cudnn86）：

   pip install --upgrade jax jaxlib==0.4.14+cuda11.cudnn86 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
   

3. 更新AlphaFold2到最新版本（官方已修复jax 0.4.x的兼容问题）：

   cd /path/to/your/alphafold/repo
   git pull
   

4. 若仍报错，临时修改代码：找到所有调用jax.xla的代码行，替换为jax._src.xla（仅作为临时 workaround，优先使用官方更新代码）。

方案3：切换到Python 3.10环境 #

Python 3.10支持jax 0.3.x系列版本，可创建独立虚拟环境：

1. 安装Python 3.10（若未安装）
2. 创建并激活虚拟环境：
   • Linux/macOS：

     python3.10 -m venv alphafold_env
     source alphafold_env/bin/activate
     

   • Windows：

     python3.10 -m venv alphafold_env
     alphafold_env\Scripts\activate
     

3. 安装指定版本的jax/jaxlib及AlphaFold2依赖：

   pip install "jax<=0.3.16" "jaxlib<=0.3.16" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
   pip install -r /path/to/your/alphafold/requirements.txt
   

4. 激活环境后运行AlphaFold2即可。

内容的提问来源于stack exchange，提问作者jfslet