安装

安装 JAX 极其简单，类似于安装 NumPy 或其他任何 Python 库。

仅 CPU 安装

如果你没有 GPU 支持，仍然可以在 CPU 上运行 JAX。

1. 升级 pip：pip install --upgrade pip
2. 安装 JAX：pip install --upgrade "jax[cpu]"

1pip install --upgrade pip
2pip install --upgrade "jax[cpu]"

---

GPU 安装

与其他深度学习框架类似，JAX 支持 CUDA，但我们需要自行安装 CUDA 和 CuDNN。

安装 CUDA

由于 CUDA 的安装依赖于多个因素（如操作系统类型、发行版、版本等），以下示例基于最常见的配置：在 x86_64 架构上使用 Ubuntu 20.04：

1wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
2sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
3wget https://developer.download.nvidia.com/compute/cuda/11.4.1/local_installers/cuda-repo-ubuntu2004-11-4-local_11.4.1-470.57.02-1_amd64.deb
4sudo dpkg -i cuda-repo-ubuntu2004-11-4-local_11.4.1-470.57.02-1_amd64.deb
5sudo apt-key add /var/cuda-repo-ubuntu2004-11-4-local/7fa2af80.pub
6sudo apt-get update
7sudo apt-get -y install cuda

对于其他架构、操作系统版本或其他配置，请参考 NVIDIA 提供的 CUDA 下载页面 获取对应安装方法。

安装 CuDNN

安装 CuDNN 需要注册 NVIDIA 开发者计划，请在下载前完成注册。

适用于 JAX 的 CUDA 10/11 版本需要搭配 CuDNN 7 或 8 版本使用。

使用 pip 安装

在终端中执行以下命令：

1pip install --upgrade pip
2pip install --upgrade "jax[cuda111]" -f https://storage.googleapis.com/jax-releases/jax_releases.html

第二条命令中 jax 后面的 [cuda111] 取决于系统中安装的 CUDA 版本。可参考下表进行选择：

使用命令 nvcc --version 可查看当前系统的 CUDA 版本。

TPU 安装

可以通过在 TPU 云虚拟机中运行以下命令来安装 JAX：

1pip install --upgrade pip
2pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

从源码构建

如果你是高级用户，并希望从源码编译安装，请先安装必要的依赖项。

前置条件

• C++ 编译器：
  sudo apt install g++
• Python 及相关开发库：
  sudo apt install python python3-dev
• 必需的 Python 库（通过 pip 安装）：
  pip install six numpy scipy wheel

安装 jaxlib

可以使用以下脚本安装包含 XLA 的 jaxlib：

1python build/build.py --enable_cuda
2pip install dist/*.whl

注意：--enable_cuda 是可选项，若仅构建 CPU 版本，可省略该参数。

安装 jax

完成 jaxlib 安装后，即可通过以下命令安装 jax：

pip install -e .

完整脚本如下，仅供参考：

1sudo apt install g++ python python3-dev
2pip install six numpy scipy wheel
3python build/build.py --enable_cuda
4pip install dist/*.whl
5pip install -e .