安装

安装#

使用 JAX 需要安装两个软件包：jax，它是纯 Python 且跨平台的，以及 jaxlib，它包含编译后的二进制文件，并且需要为不同的操作系统和加速器构建不同的版本。

摘要：对于大多数用户，典型的 JAX 安装可能如下所示

仅 CPU (Linux/macOS/Windows)
```
pip install -U jax
```
GPU (NVIDIA, CUDA 12)
```
pip install -U "jax[cuda12]"
```
TPU (Google Cloud TPU VM)
```
pip install -U "jax[tpu]"
```

支持的平台#

下表显示了所有支持的平台和安装选项。检查您的设置是否受支持；如果显示“是”或“实验性”，请单击相应的链接以了解如何更详细地安装 JAX。

	Linux, x86_64	Linux, aarch64	Mac, x86_64	Mac, aarch64	Windows, x86_64	Windows WSL2, x86_64
CPU	是	是	jax≤0.4.38 仅限	是	是	是
NVIDIA GPU	是	是	否	不适用	否	实验性
Google Cloud TPU	是	不适用	不适用	不适用	不适用	不适用
AMD GPU	实验性	否	实验性	不适用	否	否
Apple GPU	不适用	否	不适用	实验性	不适用	不适用
Intel GPU	实验性	不适用	不适用	不适用	否	否

CPU#

pip 安装：CPU#

目前，JAX 团队为以下操作系统和架构发布 jaxlib wheels

Linux, x86_64
Linux, aarch64
macOS, Intel
macOS, Apple ARM-based
Windows, x86_64 (实验性)

要安装仅 CPU 版本的 JAX，这可能对在笔记本电脑上进行本地开发很有用，您可以运行

pip install --upgrade pip
pip install --upgrade jax

在 Windows 上，如果您的机器上尚未安装 Microsoft Visual Studio 2019 Redistributable，您可能还需要安装它。

其他操作系统和架构需要从源代码构建。尝试在其他操作系统和架构上使用 pip 安装可能会导致 jaxlib 未与 jax 一起安装，尽管 jax 可能会成功安装（但在运行时会失败）。

NVIDIA GPU#

JAX 支持 SM 版本 5.2 (Maxwell) 或更新版本的 NVIDIA GPU。请注意，Kepler 系列 GPU 不再受 JAX 支持，因为 NVIDIA 已在其软件中放弃对 Kepler GPU 的支持。

您必须首先安装 NVIDIA 驱动程序。建议您安装 NVIDIA 提供的最新驱动程序，但对于 Linux 上的 CUDA 12，驱动程序版本必须 >= 525.60.13。

如果您需要在较旧的驱动程序上使用较新的 CUDA 工具包，例如在您无法轻松更新 NVIDIA 驱动程序的集群上，您或许可以使用 NVIDIA 为此目的提供的 CUDA 向前兼容性软件包。

pip 安装：NVIDIA GPU（CUDA，通过 pip 安装，更简单）#

有两种方法可以安装支持 NVIDIA GPU 的 JAX

使用从 pip wheels 安装的 NVIDIA CUDA 和 cuDNN
使用自行安装的 CUDA/cuDNN

JAX 团队强烈建议使用 pip wheels 安装 CUDA 和 cuDNN，因为它要容易得多！

NVIDIA 仅针对 x86_64 和 aarch64 发布了 CUDA pip 软件包；在其他平台上，您必须使用 CUDA 的本地安装。

pip install --upgrade pip

# NVIDIA CUDA 12 installation
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda12]"

如果 JAX 检测到错误版本的 NVIDIA CUDA 库，您需要检查以下几件事

确保未设置 LD_LIBRARY_PATH，因为 LD_LIBRARY_PATH 可能会覆盖 NVIDIA CUDA 库。
确保安装的 NVIDIA CUDA 库是 JAX 要求的那些。重新运行上面的安装命令应该有效。

pip 安装：NVIDIA GPU（CUDA，本地安装，更难）#

如果您更喜欢使用预安装的 NVIDIA CUDA 副本，则必须首先安装 NVIDIA CUDA 和 cuDNN。

JAX 仅为 Linux x86_64 和 Linux aarch64 提供预构建的 CUDA 兼容 wheels。其他操作系统和架构的组合是可能的，但需要从源代码构建（有关更多信息，请参阅从源代码构建）。

您应该使用 NVIDIA 驱动程序版本，该版本至少与您的 NVIDIA CUDA 工具包的相应驱动程序版本一样新。如果您需要在较旧的驱动程序上使用较新的 CUDA 工具包，例如在您无法轻松更新 NVIDIA 驱动程序的集群上，您或许可以使用 NVIDIA 为此目的提供的 CUDA 向前兼容性软件包。

JAX 目前发布一个 CUDA wheel 变体

构建于	兼容于
CUDA 12.3	CUDA >=12.1
CUDNN 9.1	CUDNN >=9.1, <10.0
NCCL 2.19	NCCL >=2.18

JAX 检查您的库的版本，如果它们不够新，则会报告错误。设置 JAX_SKIP_CUDA_CONSTRAINTS_CHECK 环境变量将禁用此检查，但使用旧版本的 CUDA 可能会导致错误或不正确的结果。

NCCL 是一个可选的依赖项，仅当您执行多 GPU 计算时才需要。

要安装，请运行

pip install --upgrade pip

# Installs the wheel compatible with NVIDIA CUDA 12 and cuDNN 9.0 or newer.
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda12-local]"

这些 pip 安装不适用于 Windows，并且可能会静默失败；请参阅上面的表格。

您可以使用以下命令找到您的 CUDA 版本

nvcc --version

JAX 使用 LD_LIBRARY_PATH 查找 CUDA 库，并使用 PATH 查找二进制文件（ptxas、nvlink）。请确保这些路径指向正确的 CUDA 安装。

JAX 需要 libdevice10.bc，它通常来自 cuda-nvvm 软件包。请确保它存在于您的 CUDA 安装中。

如果您在使用预构建的 wheels 时遇到任何错误或问题，请在 GitHub issue tracker 上告知 JAX 团队。

NVIDIA GPU Docker 容器#

NVIDIA 提供 JAX Toolbox 容器，这些容器是最新的容器，包含 jax 和一些模型/框架的 nightly 版本。

Google Cloud TPU#

pip 安装：Google Cloud TPU#

JAX 为 Google Cloud TPU 提供预构建的 wheels。要安装 JAX 以及适当版本的 jaxlib 和 libtpu，您可以在您的云 TPU VM 中运行以下命令

pip install "jax[tpu]"

对于 Colab (https://colab.research.google.com/) 的用户，请确保您使用的是 TPU v2，而不是较旧的、已弃用的 TPU 运行时。

Mac GPU#

pip 安装#

Apple 提供了一个实验性的 Metal 插件。有关详细信息，请参阅 Apple 的 JAX on Metal 文档。

注意：Metal 插件有几个注意事项

Metal 插件是新的且实验性的，并且存在许多已知问题。请在 JAX issue tracker 上报告任何问题。
Metal 插件目前需要非常特定版本的 jax 和 jaxlib。随着插件 API 的成熟，此限制将随着时间的推移而放宽。

AMD GPU (Linux)#

JAX 具有实验性的 ROCm 支持。有两种安装 JAX 的方法

使用 AMD 的 Docker 容器；或
从源代码构建。请参阅 Additional notes for building a ROCm jaxlib for AMD GPUs 部分。

Intel GPU#

Intel 提供了一个实验性的 OneAPI 插件：intel-extension-for-openxla，用于 Intel GPU 硬件。有关更多详细信息和安装说明，请参阅以下两种方法之一

Pip 安装：JAX 在 Intel GPU 上的加速。
使用 Intel 的 XLA Docker 容器。

请报告与以下内容相关的任何问题

JAX：JAX issue tracker。
Intel 的 OpenXLA 插件：Intel-extension-for-openxla issue tracker。

Conda (社区支持)#

Conda 安装#

有一个社区支持的 jax Conda 构建版本。要使用 conda 安装它，只需运行

conda install jax -c conda-forge

如果您在具有 NVIDIA GPU 的机器上运行此命令，则应安装 CUDA 启用的 jaxlib 软件包。

要确保您安装的 jax 版本确实是 CUDA 启用的，请运行

conda install "jaxlib=*=*cuda*" jax -c conda-forge

如果您想覆盖 JAX 使用的 CUDA 版本，或者在没有 GPU 的机器上安装 CUDA 构建版本，请按照 conda-forge 网站的 Tips & tricks 部分中的说明进行操作。

有关更多详细信息，请访问 conda-forge jaxlib 和 jax 仓库。

JAX nightly 安装#

Nightly 版本反映了构建时主 JAX 仓库的状态，并且可能无法通过完整的测试套件。

与安装 JAX 版本的说明不同，这里我们在命令行上显式命名 JAX 的所有软件包，因此如果新版本可用，pip 将升级它们。

仅 CPU

pip install -U --pre jax jaxlib -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html

Google Cloud TPU

pip install -U --pre jax jaxlib libtpu requests -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

NVIDIA GPU (CUDA 12)

pip install -U --pre jax jaxlib "jax-cuda12-plugin[with-cuda]" jax-cuda12-pjrt -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html

NVIDIA GPU (CUDA 12) legacy

以下用于单体 CUDA jaxlibs 的历史 nightly 版本。您很可能不需要这个；不会构建更多的单体 CUDA jaxlibs，并且现有的 jaxlibs 将在 2024 年 9 月到期。请使用上面的“CUDA 12”选项。

pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_cuda12_releases.html

从源代码构建 JAX#

请参阅从源代码构建。

安装旧版本的 `jaxlib` wheels#

由于 Python 包索引的存储限制，JAX 团队会定期从 https://pypi.ac.cn/project/jax 上的版本中删除旧版本的 jaxlib wheels。这些仍然可以通过此处的 URL 直接安装。例如

# Install jaxlib on CPU via the wheel archive
pip install "jax[cpu]==0.3.25" -f https://storage.googleapis.com/jax-releases/jax_releases.html

# Install the jaxlib 0.3.25 CPU wheel directly
pip install jaxlib==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_releases.html

对于特定的旧 GPU wheels，请务必使用 jax_cuda_releases.html URL；例如

pip install jaxlib==0.3.25+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

安装

目录

安装#