从源码构建#

首先，获取 JAX 源代码

git clone https://github.com/jax-ml/jax
cd jax

构建 JAX 包含两个步骤

  1. 构建或安装 jaxlib，这是 jax 的 C++ 支持库。

  2. 安装 jax Python 包。

构建或安装 jaxlib#

使用 pip 安装 jaxlib#

如果您只修改 JAX 的 Python 部分，我们建议使用 pip 从预构建的 wheel 安装 jaxlib。

pip install jaxlib

有关 pip 安装（例如，GPU 和 TPU 支持）的完整指南，请参阅 JAX README。

从源码构建 jaxlib#

警告

尽管通常可以使用大多数现代编译器从源码编译 jaxlib，但构建仅使用 clang 进行测试。欢迎提交 PR 以改进对不同工具链的支持，但其他编译器不被积极支持。

要从源码构建 jaxlib，您还必须安装一些先决条件

  • C++ 编译器

    如上方的框中所述，最好使用最新版本的 clang（撰写本文时，我们测试的版本是 18），但其他编译器（例如 g++ 或 MSVC）可能也能工作。

    在 Ubuntu 或 Debian 上，您可以按照 LLVM 文档中的说明安装最新稳定版本的 clang。

    如果您在 Mac 上构建，请确保已安装 XCode 和 XCode 命令行工具。

    请参阅下文了解 Windows 构建说明。

  • Python：用于运行构建辅助脚本。请注意，无需在本地安装 Python 依赖项，因为在构建过程中将忽略您的系统 Python；请参阅 管理 hermetic Python 获取详细信息。

要为 CPU 或 TPU 构建 jaxlib，您可以运行

python build/build.py build --wheels=jaxlib --verbose
pip install dist/*.whl  # installs jaxlib (includes XLA)

要构建与当前系统安装不同的 Python 版本的 wheel，请将 --python_version 标志传递给构建命令

python build/build.py build --wheels=jaxlib --python_version=3.12 --verbose

本文档的其余部分假定您正在为与当前系统安装匹配的 Python 版本构建。如果您需要为不同版本构建，只需在每次调用 python build/build.py 时追加 --python_version=<py version> 标志。请注意，无论是否传递 --python_version 参数，Bazel 构建始终会使用 hermetic Python 安装。

如果您想构建 jaxlib 和 CUDA 插件：运行

python build/build.py build --wheels=jaxlib,jax-cuda-plugin,jax-cuda-pjrt

生成三个 wheel（不带 cuda 的 jaxlib、jax-cuda-plugin 和 jax-cuda-pjrt）。默认情况下，所有 CUDA 编译步骤都由 NVCC 和 clang 执行，但可以通过 --build_cuda_with_clang 标志将其限制为 clang。

有关配置选项，请参阅 python build/build.py --help。此处 python 应为您的 Python 3 解释器的名称；在某些系统上，您可能需要改用 python3。尽管使用 python 调用脚本，但 Bazel 始终会使用其自己的 hermetic Python 解释器和依赖项，只有 build/build.py 脚本本身将由您的系统 Python 解释器处理。默认情况下，wheel 会写入当前目录的 dist/ 子目录。

  • JAX 版本从 v.0.4.32 开始：您可以在配置选项中提供自定义的 CUDA 和 CUDNN 版本。Bazel 将下载它们并将其用作目标依赖项。

    要下载特定版本的 CUDA/CUDNN redistributions，您可以使用 --cuda_version 和 --cudnn_version 标志

    python build/build.py build --wheels=jax-cuda-plugin --cuda_version=12.3.2 \
--cudnn_version=9.1.1

    或

    python build/build.py build --wheels=jax-cuda-pjrt --cuda_version=12.3.2 \
--cudnn_version=9.1.1

    请注意，这些参数是可选的：默认情况下，Bazel 将下载 .bazelrc 中分别在环境变量 HERMETIC_CUDA_VERSION 和 HERMETIC_CUDNN_VERSION 中提供的 CUDA 和 CUDNN redistribution 版本。

    要指向本地文件系统上的 CUDA/CUDNN/NCCL redistributions，您可以使用以下命令

    python build/build.py build --wheels=jax-cuda-plugin \
--bazel_options=--repo_env=LOCAL_CUDA_PATH="/foo/bar/nvidia/cuda" \
--bazel_options=--repo_env=LOCAL_CUDNN_PATH="/foo/bar/nvidia/cudnn" \
--bazel_options=--repo_env=LOCAL_NCCL_PATH="/foo/bar/nvidia/nccl"

    请参阅 XLA 文档 中的完整说明。

  • JAX 版本 v.0.4.32 之前：您必须安装 CUDA 和 CUDNN，并通过配置选项提供它们的路径。

使用修改后的 XLA 仓库从源码构建 jaxlib。#

JAX 依赖于 XLA，其源代码位于 XLA GitHub 仓库。默认情况下，JAX 使用 pinned 版本的 XLA 仓库，但在处理 JAX 时，我们经常希望使用本地修改过的 XLA 副本。有两种方法可以实现此目的

  • 使用 Bazel 的 override_repository 功能，您可以将其作为命令行标志传递给 build.py，如下所示

    python build/build.py build --wheels=jaxlib --local_xla_path=/path/to/xla

  • 修改 JAX 源代码树根目录下的 WORKSPACE 文件，将其指向不同的 XLA 树。

要将更改贡献回 XLA，请向 XLA 仓库提交 PR。

JAX 固定的 XLA 版本会定期更新，尤其是在每次 jaxlib 发布之前。

在 Windows 上从源码构建 jaxlib 的附加说明#

注意：JAX 不支持 Windows 上的 CUDA；请使用 WSL2 以获得 CUDA 支持。

在 Windows 上，请遵循 安装 Visual Studio 来设置 C++ 工具链。需要 Visual Studio 2019 版本 16.5 或更高版本。

JAX 构建使用符号链接，这需要您激活 开发者模式。

您可以选择使用其 Windows 安装程序 安装 Python，或者如果您愿意，可以使用 Anaconda 或 Miniconda 来设置 Python 环境。

Bazel 的某些目标使用 bash 工具进行脚本编写，因此需要 MSYS2。有关更多详细信息，请参阅 在 Windows 上安装 Bazel。安装以下软件包

pacman -S patch coreutils

安装 coreutils 后，realpath 命令应该在您的 shell 路径中。

一切安装完成后。打开 PowerShell，并确保 MSYS2 在当前会话的路径中。确保 bazel、patch 和 realpath 是可访问的。激活 conda 环境。

python .\build\build.py build --wheels=jaxlib

要生成带有调试信息的构建，请添加标志 --bazel_options='--copt=/Z7'。

为 AMD GPU 构建 ROCM jaxlib 的附加说明#

有关使用 ROCm 支持构建 jaxlib 的详细说明，请参阅官方指南：从源码构建 ROCm JAX

管理 hermetic Python#

为确保 JAX 的构建是可重现的，在支持的平台（Linux、Windows、MacOS）上行为一致，并且与本地系统的特定配置正确隔离，我们在所有通过 Bazel 执行的构建和测试命令中都依赖于 hermetic Python（由 rules_python 提供，有关详细信息，请参阅 工具链注册）。这意味着您的系统 Python 安装将在构建期间被忽略，Python 解释器本身以及所有 Python 依赖项将由 bazel 直接管理。

指定 Python 版本#

当您运行 build/build.py 工具时，hermetic Python 的版本会自动设置为匹配您用于运行 build/build.py 脚本的 Python 版本。要显式选择特定版本，您可以将 --python_version 参数传递给工具

python build/build.py build --python_version=3.12

在后台，hermetic Python 版本由 HERMETIC_PYTHON_VERSION 环境变量控制，该变量在您运行 build/build.py 时自动设置。如果您直接运行 bazel，您可能需要以以下方式之一显式设置该变量

# Either add an entry to your `.bazelrc` file
build --repo_env=HERMETIC_PYTHON_VERSION=3.12

# OR pass it directly to your specific build command
bazel build <target> --repo_env=HERMETIC_PYTHON_VERSION=3.12

# OR set the environment variable globally in your shell:
export HERMETIC_PYTHON_VERSION=3.12

您可以通过在运行之间切换 --python_version 的值，在同一台机器上依次针对不同版本的 Python 运行构建和测试。构建的所有与 Python 无关的部分将从之前的构建中缓存并重用，供后续构建使用。

指定 Python 依赖项#

在 bazel 构建期间，所有 JAX 的 Python 依赖项都被固定到其特定版本。这是确保构建可重现性所必需的。JAX 依赖项的完整传递闭包及其相应哈希的固定版本在 build/requirements_lock_<python version>.txt 文件中指定（例如，对于 Python 3.12，为 build/requirements_lock_3_12.txt）。

要更新锁定文件，请确保 build/requirements.in 包含所需的直接依赖项列表，然后执行以下命令（该命令将在后台调用 pip-compile）

python build/build.py requirements_update --python_version=3.12

或者，如果您需要更多控制，可以 B直接运行 bazel 命令（这两个命令等效）

# Regular Python
bazel run //build:requirements.update --repo_env=HERMETIC_PYTHON_VERSION=3.12

# Free-threaded Python
bazel run //build:requirements_ft.update --repo_env=HERMETIC_PYTHON_VERSION=3.13-ft

请注意，由于底层仍然使用 pip 和 pip-compile 工具，因此这些工具支持的大部分命令行参数和功能也将被 Bazel 需求更新程序命令识别。例如，如果您希望更新程序考虑预发布版本，只需将 --pre 参数传递给 bazel 命令

bazel run //build:requirements.update --repo_env=HERMETIC_PYTHON_VERSION=3.12 -- --pre

指定对本地 wheel 的依赖#

默认情况下，构建会扫描仓库根目录下的 dist 目录以查找任何要包含在依赖项列表中的本地 .whl 文件。如果 wheel 特定于 Python 版本，则仅包含与所选 Python 版本匹配的 wheel。

整体本地 wheel 搜索和选择逻辑由 WORKSPACE 文件中 python_init_repositories() 宏的参数控制（直接从 WORKSPACE 文件调用）。您可以使用 local_wheel_dist_folder 来更改本地 wheel 文件夹的位置。使用 local_wheel_inclusion_list 和 local_wheel_exclusion_list 参数指定在搜索中应包含和/或排除的 wheel（支持基本通配符匹配）。

如果需要，您还可以手动依赖本地 .whl 文件，绕过自动本地 wheel 搜索机制。例如，要依赖您新构建的 jaxlib wheel，您可以在 build/requirements.in 中添加 wheel 的路径，然后为选定的 Python 版本重新运行需求更新程序命令。例如

echo -e "\n$(realpath jaxlib-0.4.27.dev20240416-cp312-cp312-manylinux_2_27_x86_64.whl)" >> build/requirements.in
python build/build.py requirements_update --python_version=3.12

指定对 nightly wheel 的依赖#

为了构建和测试最新的、可能不稳定的 Python 依赖集，我们提供了一个特殊的依赖项更新程序命令，如下所示

python build/build.py requirements_update --python_version=3.12 --nightly_update

或者，如果您直接运行 bazel（这两个命令等效）

bazel run //build:requirements_nightly.update --repo_env=HERMETIC_PYTHON_VERSION=3.12

它与常规更新程序之间的区别是，默认情况下它会接受预发布、开发版和 nightly 包，它还会搜索 https://pypi.anaconda.org/scientific-python-nightly-wheels/simple 作为额外的索引 URL，并且不会将哈希值放入最终的需求锁定文件中。

定制 hermetic Python (高级用法)#

我们开箱即用地支持所有当前版本的 Python，因此除非您的工作流程有非常特殊的要求（例如使用自己的自定义 Python 解释器的能力），否则您可以安全地跳过本节。

简而言之，如果您依赖于非标准的 Python 工作流程，您仍然可以实现 hermetic Python 设置的高度灵活性。概念上，与非 hermetic 情况相比，只有一个区别：您需要以文件的形式进行思考，而不是以安装的形式进行思考（即，思考您的构建实际上依赖于哪些文件，而不是哪些文件需要安装在您的系统上），其余部分基本相同。

因此，实际上，要完全控制您的 Python 环境（无论是 hermetic 还是非 hermetic），您需要能够执行以下三项操作

  1. 指定要使用的 Python 解释器（即，选择实际的 python 或 python3 二进制文件以及同目录下的相关库）。

  2. 指定 Python 依赖项列表（例如 numpy）及其实际版本。

  3. 能够轻松地添加/删除/更新列表中的依赖项。每个依赖项本身也可以是自定义的（例如自构建的）。

您已经知道如何在非 hermetic Python 环境中执行所有这些步骤，以下是在 hermetic 环境中执行相同操作的方法（通过以文件为中心的方法，而不是以安装为中心的方法）

  1. 不要安装 Python，而是将 Python 解释器打包成 tar 或 zip 文件。根据您的情况，您可以直接获取许多现有的文件之一（例如 python-build-standalone），或者构建自己的文件并将其打包成存档（遵循官方 构建说明 即可）。例如，在 Linux 上，它将如下所示

    ./configure --prefix python
make -j12
make altinstall
tar -czpf my_python.tgz python

    准备好 tarball 后，通过将 HERMETIC_PYTHON_URL 环境变量指向存档（本地或互联网上的），将其集成到构建中

    --repo_env=HERMETIC_PYTHON_URL="file:///local/path/to/my_python.tgz"
--repo_env=HERMETIC_PYTHON_SHA256=<file's_sha256_sum>

# OR
--repo_env=HERMETIC_PYTHON_URL="https://remote/url/to/my_python.tgz"
--repo_env=HERMETIC_PYTHON_SHA256=<file's_sha256_sum>

# We assume that top-level folder in the tarball is called "python", if it is
# something different just pass additional HERMETIC_PYTHON_PREFIX parameter
--repo_env=HERMETIC_PYTHON_URL="https://remote/url/to/my_python.tgz"
--repo_env=HERMETIC_PYTHON_SHA256=<file's_sha256_sum>
--repo_env=HERMETIC_PYTHON_PREFIX="my_python/install"

  2. 不要执行 pip install，而是创建 requirements_lock.txt 文件，其中包含您依赖项的完整传递闭包。您也可以依赖此仓库中已有的依赖项（只要它们与您的自定义 Python 版本兼容）。没有关于如何执行此操作的特殊说明，您可以遵循本文档 指定 Python 依赖项 中推荐的步骤，直接调用 pip-compile（请注意，锁定文件必须是 hermetic 的，但您可以随时从非 hermetic Python 生成它）甚至手动创建它（请注意，哈希值在锁定文件中是可选的）。

  3. 如果您需要更新或自定义依赖项列表，您可以再次遵循 指定 Python 依赖项 的说明来更新 requirements_lock.txt，直接调用 pip-compile 或手动修改它。如果您有要使用的自定义软件包，只需在锁定文件中直接指向其 .whl 文件（请记住，以文件为中心，而不是以安装为中心）（请注意，requirements.txt 和 requirements_lock.txt 文件支持本地 wheel 引用）。如果您的 requirements_lock.txt 已在 WORKSPACE 文件中指定为 python_init_repositories() 的依赖项，则无需再做任何其他操作。否则，您可以指向您的自定义文件，如下所示

    --repo_env=HERMETIC_REQUIREMENTS_LOCK="/absolute/path/to/custom_requirements_lock.txt"

    另外请注意，如果您使用 HERMETIC_REQUIREMENTS_LOCK，那么它将完全控制您的依赖项列表，并且 指定本地 wheel 依赖项 中描述的自动本地 wheel 解析逻辑将被禁用，以免与其发生冲突。

就这样。总结一下：如果您有一个包含 Python 解释器的存档和一个包含依赖项完整传递闭包的 requirements_lock.txt 文件，那么您就完全控制了您的 Python 环境。

自定义 hermetic Python 示例#

注意，对于下面的所有示例，您也可以全局设置环境变量（即在 shell 中 export 而不是作为 --repo_env 参数传递给命令），因此通过 build/build.py 调用 bazel 即可正常工作。

从互联网构建自定义 Python 3.13，使用此仓库中已包含的默认 requirements_lock_3_13.txt（即自定义解释器但默认依赖项）

bazel build <target>
  --repo_env=HERMETIC_PYTHON_VERSION=3.13
  --repo_env=HERMETIC_PYTHON_URL="https://github.com/indygreg/python-build-standalone/releases/download/20241016/cpython-3.13.0+20241016-x86_64-unknown-linux-gnu-install_only.tar.gz"
  --repo_env=HERMETIC_PYTHON_SHA256="2c8cb15c6a2caadaa98af51df6fe78a8155b8471cb3dd7b9836038e0d3657fb4"

从本地文件系统构建自定义 Python 3.13 和自定义锁定文件（假设在运行命令之前，锁定文件已放入此仓库的 jax/build 文件夹中）

bazel test <target>
  --repo_env=HERMETIC_PYTHON_VERSION=3.13
  --repo_env=HERMETIC_PYTHON_URL="file:///path/to/cpython.tar.gz"
  --repo_env=HERMETIC_PYTHON_PREFIX="prefix/to/strip/in/cython/tar/gz/archive"
  --repo_env=HERMETIC_PYTHON_SHA256=<sha256_sum>
  --repo_env=HERMETIC_REQUIREMENTS_LOCK="/absolute/path/to/build:custom_requirements_lock.txt"

如果默认 Python 解释器对您来说足够好，而您只需要一组自定义的依赖项

bazel test <target>
  --repo_env=HERMETIC_PYTHON_VERSION=3.13
  --repo_env=HERMETIC_REQUIREMENTS_LOCK="/absolute/path/to/build:custom_requirements_lock.txt"

注意，您可以为同一 Python 版本拥有多个不同的 requirement_lock.txt 文件，以支持不同的场景。您可以通过在 WORKSPACE 文件中指定 HERMETIC_PYTHON_VERSION 来控制选择哪个文件。

requirements = {
  "3.11": "//build:requirements_lock_3_11.txt",
  "3.12": "//build:requirements_lock_3_12.txt",
  "3.13": "//build:requirements_lock_3_13.txt",
  "3.13-scenario1": "//build:scenario1_requirements_lock_3_13.txt",
  "3.13-scenario2": "//build:scenario2_requirements_lock_3_13.txt",
},

然后，您可以在不更改环境中的任何内容的情况下，构建和测试各种组合。

# To build with scenario1 dependencies:
bazel test <target> --repo_env=HERMETIC_PYTHON_VERSION=3.13-scenario1

# To build with scenario2 dependencies:
bazel test <target> --repo_env=HERMETIC_PYTHON_VERSION=3.13-scenario2

# To build with default dependencies:
bazel test <target> --repo_env=HERMETIC_PYTHON_VERSION=3.13

# To build with scenario1 dependencies and custom Python 3.13 interpreter:
bazel test <target>
  --repo_env=HERMETIC_PYTHON_VERSION=3.13-scenario1
  --repo_env=HERMETIC_PYTHON_URL="file:///path/to/cpython.tar.gz"
  --repo_env=HERMETIC_PYTHON_SHA256=<sha256_sum>

安装 jax#

一旦 jaxlib 安装完毕，您就可以通过运行以下命令来安装 jax

pip install -e .  # installs jax

要升级到 GitHub 上的最新版本，只需从 JAX 仓库根目录运行 git pull，然后通过运行 build.py 进行重建，或者在必要时升级 jaxlib。您不应该需要重新安装 jax，因为 pip install -e 会设置从 site-packages 到仓库的符号链接。

运行测试#

运行 JAX 测试有两种支持的机制：使用 Bazel 或使用 pytest。

使用 Bazel#

首先，使用 --configure_only 标志配置 JAX 构建。对于 CPU 测试，请传递 --wheel_list=jaxlib；对于 GPU 测试，请传递 CUDA/ROCM。

python build/build.py build --wheels=jaxlib --configure_only
python build/build.py build --wheels=jax-cuda-plugin --configure_only
python build/build.py build --wheels=jax-rocm-plugin --configure_only

您可以将附加选项传递给 build.py 来配置构建；有关详细信息，请参阅 jaxlib 构建文档。

默认情况下，Bazel 构建使用从源码构建的 jaxlib 来运行 JAX 测试。要运行 JAX 测试，请运行

bazel test //tests:cpu_tests //tests:backend_independent_tests

如果您拥有必要的硬件，也可以使用 //tests:gpu_tests 和 //tests:tpu_tests。

您需要配置 cuda 来运行 gpu 测试。

python build/build.py build --wheels=jaxlib,jax-cuda-plugin,jax-cuda-pjrt --configure_only

要使用预安装的 jaxlib 而不是从源码构建，您需要先将其在 hermetic Python 中可用。要为 hermetic Python 安装特定版本的 jaxlib（以 jaxlib >= 0.4.26 为例），请运行

echo -e "\njaxlib >= 0.4.26" >> build/requirements.in
python build/build.py requirements_update

或者，要从本地 wheel 安装 jaxlib（假设使用 Python 3.12）

echo -e "\n$(realpath jaxlib-0.4.26-cp312-cp312-manylinux_2_27_x86_64.whl)" >> build/requirements.in
python build/build.py requirements_update --python_version=3.12

一旦在 hermetic 环境中安装了 jaxlib，请运行

bazel test --//jax:build_jaxlib=false //tests:cpu_tests //tests:backend_independent_tests

可以通过环境变量（见下文）控制许多测试行为。可以使用 Bazel 的 --test_env=FLAG=value 标志将环境变量传递给 JAX 测试。

JAX 的一些测试是针对多个加速器的（例如 GPU、TPU）。当 JAX 已安装时，您可以通过以下方式运行 GPU 测试

bazel test //tests:gpu_tests --local_test_jobs=4 --test_tag_filters=multiaccelerator --//jax:build_jaxlib=false --test_env=XLA_PYTHON_CLIENT_ALLOCATOR=platform

您可以使单个加速器测试在多个加速器上并行运行以加快速度。这也会触发每个加速器上的多个并发测试。对于 GPU，您可以这样做

NB_GPUS=2
JOBS_PER_ACC=4
J=$((NB_GPUS * JOBS_PER_ACC))
MULTI_GPU="--run_under $PWD/build/parallel_accelerator_execute.sh --test_env=JAX_ACCELERATOR_COUNT=${NB_GPUS} --test_env=JAX_TESTS_PER_ACCELERATOR=${JOBS_PER_ACC} --local_test_jobs=$J"
bazel test //tests:gpu_tests //tests:backend_independent_tests --test_env=XLA_PYTHON_CLIENT_PREALLOCATE=false --test_tag_filters=-multiaccelerator $MULTI_GPU

使用 pytest#

首先，通过运行 pip install -r build/test-requirements.txt 安装依赖项。

为了使用 pytest 运行所有 JAX 测试，我们建议使用 pytest-xdist，它可以并行运行测试。它已作为 pip install -r build/test-requirements.txt 命令的一部分安装。

从仓库根目录运行

pytest -n auto tests

控制测试行为#

JAX 以组合方式生成测试用例，您可以使用 JAX_NUM_GENERATED_CASES 环境变量来控制为每个测试生成的用例数量（默认为 10）。自动测试目前默认使用 25。

例如，可以这样写

# Bazel
bazel test //tests/... --test_env=JAX_NUM_GENERATED_CASES=25`

或

# pytest
JAX_NUM_GENERATED_CASES=25 pytest -n auto tests

自动测试还使用默认的 64 位浮点数和整数运行测试（JAX_ENABLE_X64）

JAX_ENABLE_X64=1 JAX_NUM_GENERATED_CASES=25 pytest -n auto tests

您可以使用 pytest 内置的选择机制运行更特定的测试集，或者您可以直接运行单个测试文件以查看有关正在运行的用例的更详细信息。

JAX_NUM_GENERATED_CASES=5 python tests/lax_numpy_test.py

通过传递环境变量 JAX_SKIP_SLOW_TESTS=1，可以跳过一些已知运行缓慢的测试。

要指定要从测试文件中运行的特定测试集，可以通过 --test_targets 标志传递一个字符串或正则表达式。例如，您可以使用以下命令运行 jax.numpy.pad 的所有测试

python tests/lax_numpy_test.py --test_targets="testPad"

Colab Notebooks 在文档构建过程中会进行错误检查。

Hypothesis 测试#

一些测试使用了 hypothesis。通常，hypothesis 会使用多个示例输入进行测试，并在测试失败时尝试找到一个较小的示例，该示例仍然会导致失败：查看测试失败，查找类似下面的行，并添加消息中提到的装饰器。

You can reproduce this example by temporarily adding @reproduce_failure('6.97.4', b'AXicY2DAAAAAEwAB') as a decorator on your test case

对于交互式开发，您可以设置环境变量 JAX_HYPOTHESIS_PROFILE=interactive（或等效的标志 --jax_hypothesis_profile=interactive），将示例数量设置为 1，并跳过示例最小化阶段。

Doctests#

JAX 使用 doctest 模式的 pytest 来测试文档中的代码示例。您可以在 ci-build.yaml 中找到运行 doctests 的最新命令。例如，您可以运行

JAX_TRACEBACK_FILTERING=off XLA_FLAGS=--xla_force_host_platform_device_count=8 pytest -n auto --tb=short --doctest-glob='*.md' --doctest-glob='*.rst' docs --doctest-continue-on-failure --ignore=docs/multi_process.md

此外，JAX 在 doctest-modules 模式下运行 pytest，以确保函数文档字符串中的代码示例能够正确运行。例如，您可以在本地使用以下命令运行此操作

JAX_TRACEBACK_FILTERING=off XLA_FLAGS=--xla_force_host_platform_device_count=8 pytest --doctest-modules jax/_src/numpy/lax_numpy.py

类型检查#

我们使用 mypy 来检查类型提示。要使用与 github CI 检查相同的配置运行 mypy，您可以使用 pre-commit 框架

pip install pre-commit
pre-commit run mypy --all-files

由于 mypy 在检查所有文件时可能有点慢，因此只检查您修改过的文件可能会更方便。要做到这一点，首先暂存更改（即 `git add` 更改的文件），然后提交更改之前运行此命令。

pre-commit run mypy

代码风格检查#

JAX 使用 ruff linter 来确保代码质量。要使用与 github CI 检查相同的配置运行 ruff，您可以使用 pre-commit 框架

pip install pre-commit
pre-commit run ruff --all-files

更新文档#

要重新构建文档，请安装几个软件包

pip install -r docs/requirements.txt

然后运行

sphinx-build -b html docs docs/build/html -j auto

这可能需要很长时间，因为它会执行文档源中的许多 Notebooks；如果您希望在不执行 Notebooks 的情况下构建文档，可以运行

sphinx-build -b html -D nb_execution_mode=off docs docs/build/html -j auto

然后您可以在 docs/build/html/index.html 中查看生成的文档。

该 -j auto 选项控制构建的并行度。您可以使用一个数字替换 auto 来控制使用多少 CPU 核心。

更新 Notebooks#

我们使用 jupytext 来维护 docs/notebooks 中 Notebooks 的两个同步副本：一个为 ipynb 格式，另一个为 md 格式。前者可以直接在 Colab 中打开和执行；后者则更容易在版本控制中跟踪 diffs。

编辑 ipynb#

对于进行大量修改代码和输出的大型更改，最简单的方法是在 Jupyter 或 Colab 中编辑 Notebooks。要在 Colab 界面中编辑 Notebooks，请打开 http://colab.research.google.com 并从本地仓库 上传。根据需要进行更新，运行 所有 单元格，然后 下载 ipynb。您可能需要测试它是否正常执行，方法如上所述使用 sphinx-build。

编辑 md#

对于对 Notebooks 的文本内容进行较小更改，最简单的方法是使用文本编辑器编辑 .md 版本。

同步 Notebooks#

编辑 Notebooks 的 ipynb 或 md 版本后，您可以通过运行 jupytext --sync 来同步两个版本；例如

pip install jupytext==1.16.4
jupytext --sync docs/notebooks/thinking_in_jax.ipynb

jupytext 版本应与 .pre-commit-config.yaml 中指定的版本匹配。

要检查 markdown 和 ipynb 文件是否正确同步，您可以使用 pre-commit 框架执行与 github CI 相同的检查。

pip install pre-commit
pre-commit run jupytext --all-files

创建新的 Notebooks#

如果您正在将新的 Notebook 添加到文档中，并希望使用此处讨论的 jupytext --sync 命令，则可以使用以下命令为 Notebook 设置 jupytext

jupytext --set-formats ipynb,md:myst path/to/the/notebook.ipynb

这通过向 Notebook 文件添加一个指定所需格式的 "jupytext" 元数据字段来实现，当调用 jupytext --sync 命令时，该命令会识别该字段。

Sphinx 构建中的 Notebooks#

一些 Notebooks 会在提交前检查以及 Read the docs 构建过程中自动构建。如果单元格引发错误，构建将失败。如果错误是故意的，您可以捕获它们，或用 raises-exceptions 元数据标记单元格（示例 PR）。您必须在 .ipynb 文件中手动添加此元数据。当其他人重新保存 Notebook 时，它将被保留。

我们排除了某些 Notebooks 的构建，例如，因为它们包含长时间的计算。请参阅 docs/conf.py 中的 exclude_patterns。

在 readthedocs.io 上构建文档#

JAX 的自动生成文档位于 https://jax.net.cn/。

整个项目的文档构建由 readthedocs JAX 设置 控制。当前设置会在代码推送到 GitHub main 分支时触发文档构建。对于每个代码版本，构建过程由 .readthedocs.yml 和 docs/conf.py 配置文件驱动。

对于每次自动文档构建，您可以查看 文档构建日志。

如果您想在 Readthedocs 上测试文档生成，可以将“documentation”GitHub 标签添加到您的 PR。然后，您将能够从“docs/readthedocs.org:jax”GitHub 检查的链接中查看文档。

对于本地测试，我能够在一个新的目录中通过重放 Readthedocs 日志中看到的命令来完成

mkvirtualenv jax-docs  # A new virtualenv
mkdir jax-docs  # A new directory
cd jax-docs
git clone --no-single-branch --depth 50 https://github.com/jax-ml/jax
cd jax
git checkout --force origin/test-docs
git clean -d -f -f
workon jax-docs

python -m pip install --upgrade --no-cache-dir pip
python -m pip install --upgrade --no-cache-dir -I Pygments==2.3.1 setuptools==41.0.1 docutils==0.14 mock==1.0.1 pillow==5.4.1 alabaster>=0.7,<0.8,!=0.7.5 commonmark==0.8.1 recommonmark==0.5.0 'sphinx<2' 'sphinx-rtd-theme<0.5' 'readthedocs-sphinx-ext<1.1'
python -m pip install --exists-action=w --no-cache-dir -r docs/requirements.txt
cd docs
python `which sphinx-build` -T -E -b html -d _build/doctrees-readthedocs -D language=en . _build/html