无法从Git仓库安装jax==0.4.5及对应jaxlib的问题求助

问题根源 #

jaxlib的预编译包在PyPI上的分发受系统架构、CUDA/ROCm版本等限制，直接指定jaxlib==0.4.5可能因为你的环境没有对应预编译包导致安装失败。另外jax与jaxlib版本必须严格对齐，单独安装jax必然引发依赖冲突。

解决步骤 #

方案1：安装带环境标识的jaxlib预编译包 #

jaxlib在PyPI上的包名会附带系统、CUDA版本等后缀，比如Linux x86+CUDA 11.8对应的包是jaxlib==0.4.5+cuda11.cudnn86。你需要根据自身环境选择正确后缀：

• 先确认你的操作系统（Linux/macOS/Windows）、硬件类型（CPU/GPU）、CUDA版本（GPU用户）
• 执行对应命令：
  • 通用CPU版本：

    pip install jax==0.4.5 jaxlib==0.4.5
    

    （如果此命令失败，说明该版本的CPU预编译包已从PyPI移除，直接切换方案2）
  • Linux + CUDA 11.7：

    pip install jax==0.4.5 jaxlib==0.4.5+cuda11.cudnn85
    

  • Linux + CUDA 11.8：

    pip install jax==0.4.5 jaxlib==0.4.5+cuda11.cudnn86
    

  注：完整的版本后缀可通过jaxlib的PyPI历史版本列表查询对应环境的匹配项。

方案2：从源码编译jaxlib #

jaxlib的源码与jax同属google/jax仓库，无需单独寻找。编译步骤如下：

1. 安装编译依赖：

   pip install numpy scipy wheel cmake
   

2. 克隆仓库并切换到目标版本标签：

   git clone https://github.com/google/jax.git
   cd jax
   git checkout jax-v0.4.5
   

3. 编译并安装jaxlib：
   • CPU版本：

     python build/build.py --enable_cuda=False
     pip install dist/jaxlib-0.4.5-*.whl
     

   • GPU版本（需提前安装CUDA和CUDNN）：

     python build/build.py --enable_cuda=True
     pip install dist/jaxlib-0.4.5-*.whl
     

4. 安装对应版本的jax：

   pip install jax==0.4.5
   

方案3：通过conda安装（conda环境用户） #

conda-forge频道可能提供适配的预编译包：

conda install -c conda-forge jax=0.4.5 jaxlib=0.4.5

验证安装 #

执行以下代码确认版本一致：

import jax
import jaxlib
print(jax.__version__)
print(jaxlib.__version__)

输出均应为0.4.5

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