从源代码构建

从源代码构建#

首先，获取 JAX 源代码：

git clone https://github.com/google/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 自述文件。

从源码构建 jaxlib#

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

  • 一个 C++ 编译器 (g++、clang 或 MSVC)

    在 Ubuntu 或 Debian 上，您可以使用以下命令安装必要的先决条件：

    sudo apt install g++ python python3-dev

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

    请参阅以下Windows构建说明。

  • 无需在本地安装Python依赖项，因为在构建过程中会忽略系统的Python；详情请查看管理封闭的Python。

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

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

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

python build/build.py --python_version=3.12

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

有两种方法可以构建支持 CUDA 的 jaxlib：(1) 使用 python build/build.py --enable_cuda 生成一个支持 cuda 的 jaxlib 轮子，或 (2) 使用 python build/build.py --enable_cuda --build_gpu_plugin --gpu_plugin_cuda_version=12 生成三个轮子（不带 cuda 的 jaxlib，jax-cuda-plugin 和 jax-cuda-pjrt）。

查看 python build/build.py --help 以获取配置选项。这里的 python 应为您的 Python 3 解释器的名称；在某些系统上，您可能需要使用 python3 代替。尽管使用 python 调用脚本，Bazel 将始终使用其自己的独立 Python 解释器和依赖项，只有 build/build.py 脚本本身将由您的系统 Python 解释器处理。默认情况下，wheel 文件将写入当前目录的 dist/ 子目录。

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

    要下载特定版本的 CUDA/CUDNN 再分发包，可以使用以下命令：

    python build/build.py --enable_cuda \
--cuda_version=12.3.2 --cudnn_version=9.1.1

    要在本地文件系统上指向 CUDA/CUDNN/NCCL 的重新分发，可以使用以下命令：

    python build/build.py --enable_cuda \
--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"

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

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

JAX 依赖于 XLA，其源代码位于 XLA GitHub 仓库。默认情况下，JAX 使用固定版本的 XLA 仓库，但在开发 JAX 时，我们通常希望使用本地修改过的 XLA 副本。有两种方法可以实现这一点：

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

    python build/build.py --bazel_options=--override_repository=xla=/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

要使用调试信息进行构建，请添加标志 --bazel_options='--copt=/Z7'。

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

你需要安装几个 ROCM/HIP 库来为 ROCM 构建。例如，在安装了 AMD 的 apt 仓库 的 Ubuntu 机器上，你需要安装多个软件包：

sudo apt install miopen-hip hipfft-dev rocrand-dev hipsparse-dev hipsolver-dev \
    rccl-dev rccl hip-dev rocfft-dev roctracer-dev hipblas-dev rocm-device-libs

要使用 ROCM 支持构建 jaxlib，您可以运行以下构建命令，根据您的路径和 ROCM 版本进行适当调整。

python build/build.py --enable_rocm --rocm_path=/opt/rocm-5.7.0

AMD 对 XLA 仓库的分支可能包含上游 XLA 仓库中没有的修复。如果你在上游仓库中遇到问题，可以尝试 AMD 的分支，通过克隆他们的仓库：

git clone https://github.com/ROCmSoftwarePlatform/xla.git

并覆盖与JAX构建的XLA存储库：

python build/build.py --enable_rocm --rocm_path=/opt/rocm-5.7.0 \
  --bazel_options=--override_repository=xla=/path/to/xla-rocm

管理封闭的Python#

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

指定 Python 版本#

当你运行 build/build.py 工具时，hermetic Python 的版本会自动设置为与你运行 build/build.py 脚本时使用的 Python 版本相匹配。要明确选择特定版本，你可以向工具传递 --python_version 参数：

python build/build.py --python_version=3.12

在底层，密封的 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

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

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

其中 3.12 是你希望更新的 Python 版本。

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

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

指定对本地轮子的依赖#

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

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

如有必要，您也可以手动依赖本地的 .whl 文件，绕过自动本地轮搜索机制。例如，要依赖您新构建的 jaxlib 轮，您可以在 build/requirements.in 中添加轮的路径，并重新运行所选 Python 版本的依赖项更新命令。例如：

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

指定对夜间轮子的依赖#

要针对最新的、可能不稳定的Python依赖集进行构建和测试，我们提供了一个特殊的依赖更新器命令，如下所示：

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

或者，如果你直接运行 bazel（这两个命令是等效的）：

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

这个与常规更新器的区别在于，默认情况下它会接受预发布、开发和夜间包，它还会将 https://pypi.anaconda.org/scientific-python-nightly-wheels/simple 作为额外的索引URL进行搜索，并且不会在生成的需求锁定文件中放置哈希值。

使用预发布 Python 版本进行构建#

我们支持所有当前版本的 Python 开箱即用，但如果您需要针对不同版本进行构建和测试（例如，最新的不稳定版本，该版本尚未正式发布），请按照以下说明操作。

  1. 确保你已经安装了构建 Python 解释器本身以及关键包（如 numpy 或 scipy）所需的必要 Linux 包。在一个典型的 Debian 系统上，你可能需要安装以下包：

sudo apt-get update
sudo apt-get build-dep python3 -y
sudo apt-get install pkg-config zlib1g-dev libssl-dev -y
# to  build scipy
sudo apt-get install libopenblas-dev -y

  1. 检查你的 WORKSPACE 文件，确保它包含 custom_python_interpreter() 条目，指向你想要构建的 Python 版本。

  2. 运行 bazel build @python_dev//:python_dev -repo_env=HERMETIC_PYTHON_VERSION=3.12 来构建 Python 解释器。注意，很容易混淆用于构建的 Python 版本（出于技术原因需要，由 HERMETIC_PYTHON_VERSION=3.12 定义）和你正在构建的 Python 版本（由你在第 2 步中指定的 custom_python_interpreter() 定义）。为了构建成功，请确保你选择的用于构建的 hermetic Python 已经存在于你的配置中（实际版本无关紧要，只要它是可用的）。默认情况下，Python 二进制文件将使用 GCC 编译器构建。如果你想使用 clang 构建它，你需要设置相应的环境变量（例如 --repo_env=CC=/usr/lib/llvm-17/bin/clang --repo_env=CXX=/usr/lib/llvm-17/bin/clang++）。

  3. 检查上一个命令的输出。在它的最后部分，你将找到一个包含你新构建的Python的 python_register_toolchains() 条目的代码片段。将该代码片段复制到你的 WORKSPACE 文件中，可以紧接在 python_init_toolchains() 条目之后（以添加新的Python版本），或者替换它（以替换现有版本，例如用自定义构建的 3.12 变体替换 3.12）。代码片段是根据你的实际设置生成的，因此它应该可以直接使用，但如果你愿意，也可以自定义它（例如，更改Python的 .tgz 文件的位置，以便可以从远程下载，而不是在本地机器上）。

  4. 确保在你的 WORKSPACE 文件中，python_init_repositories() 的 requirements 参数中包含你的 Python 版本的条目。例如，对于 Python 3.13，它应该包含类似 "3.13": "//build:requirements_lock_3_13.txt" 的内容。注意，requirements 参数中的键必须始终采用 "major.minor" 版本格式，因此即使你正在构建 Python 版本 3.13.0rc1，相应的 requirements 条目也必须是 "3.13": "//build:requirements_lock_3_13.txt"，而不是 "3.13.0rc1": "//build:requirements_lock_3_13_0rc1.txt"。

  5. 对于不稳定的 Python 版本，可选地（但强烈推荐）运行 bazel build //build:all_py_deps --repo_env=HERMETIC_PYTHON_VERSION="3.13"，其中 3.13 是你在第3步中构建的 Python 解释器版本。这将使 pip 从源码拉取并构建（对于那些还没有发布二进制文件的包，例如 numpy、scipy、matplotlib、zstandard）所有 JAX 的 Python 依赖项。建议首先执行此步骤（即独立于实际的 JAX 构建），以避免在构建 JAX 本身和构建其 Python 依赖项之间发生冲突。例如，我们通常使用 clang 构建 JAX，但使用 clang 从源码构建 matplotlib 会因 GCC 和 clang 之间 LTO 行为（链接时间优化，由 -flto 标志触发）的差异而失败，并且 matplotlib 默认假设使用 GCC。如果你构建的是稳定的 Python 版本，或者通常情况下你不期望任何 Python 依赖项从源码构建（即相应 Python 版本的二进制分发已经在仓库中存在），则不需要此步骤。

  6. 恭喜，您已经构建并配置了您的自定义 JAX 项目！现在您可以像往常一样执行构建/测试命令，只需确保 HERMETIC_PYTHON_VERSION 环境变量已设置并指向您的新版本。

  7. 注意，如果你正在构建一个 Python 的预发布版本，使用你新构建的 Python 更新 requirements_lock_<python_version>.txt 文件可能会失败，因为包仓库中不会有匹配的二进制包。当没有可用的二进制包时，pip-compile 会尝试从源码构建它们，这很可能会失败，因为它比在 pip 安装过程中做同样的事情更为严格。更新不稳定版本的 Python 的需求锁定文件的推荐方法是，先为最新的稳定版本（例如 3.12）更新需求，不带哈希值（因此使用特殊的 //build:requirements_dev.update 目标），然后将结果复制到不稳定 Python 的锁定文件中（例如 3.13）：

bazel run //build:requirements_dev.update --repo_env=HERMETIC_PYTHON_VERSION="3.12"
cp build/requirements_lock_3_12.txt build/requirements_lock_3_13.txt
bazel build //build:all_py_deps --repo_env=HERMETIC_PYTHON_VERSION="3.13"
# You may need to edit manually the resultant lock file, depending on how ready
# your dependencies are for the new version of Python.

安装 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#

首先，通过运行以下命令配置 JAX 构建：

python build/build.py --configure_only

你可以向 build.py 传递额外的选项来配置构建；详情请参阅 jaxlib 构建文档。

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

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

//tests:gpu_tests 和 //tests:tpu_tests 也是可用的，如果你有必要的硬件。

要使用预安装的 jaxlib 而不是构建它，首先需要将其在封闭的 Python 环境中可用。要在封闭的 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-manylinux2014_x86_64.whl)" >> build/requirements.in
python build/build.py --requirements_update --python_version=3.12

一旦你已经 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 笔记本作为文档构建的一部分进行错误测试。

假设检验#

一些测试使用了 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，并跳过示例最小化阶段。

文档测试#

JAX 使用 pytest 的 doctest 模式来测试文档中的代码示例。你可以使用以下命令运行：

pytest docs

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

pytest --doctest-modules jax/_src/numpy/lax_numpy.py

请记住，在运行 doctest 命令时，有几个文件被标记为跳过整个包；您可以在 ci-build.yaml 中查看详细信息。

类型检查#

我们使用 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 代码检查工具来确保代码质量。要使用与 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

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

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

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

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

更新笔记本#

我们使用 jupytext 来维护 docs/notebooks 中的两个同步的笔记本副本：一个在 ipynb 格式，另一个在 md 格式。前者的优点是它可以直接在 Colab 中打开和执行；后者的优点是它使得在版本控制中跟踪差异变得更加容易。

编辑 ipynb#

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

编辑 md#

对于笔记本文本内容的小改动，最简单的方法是使用文本编辑器编辑 .md 版本。

同步笔记本#

在编辑了笔记本的 ipynb 或 md 版本后，你可以使用 jupytext 通过在更新的笔记本上运行 jupytext --sync 来同步这两个版本；例如：

pip install jupytext==1.16.0
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

创建新笔记本#

如果你正在向文档中添加一个新的笔记本，并且希望使用这里讨论的 jupytext --sync 命令，你可以通过以下命令为 jupytext 设置你的笔记本：

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

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

Sphinx 构建中的笔记本#

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

我们从构建中排除一些笔记本，例如，因为它们包含长时间的计算。请参阅 conf.py 中的 exclude_patterns。

在 readthedocs.io 上构建文档#

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

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

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

如果你想在 Readthedocs 上测试文档生成，你可以将代码推送到 test-docs 分支。该分支也会自动构建，你可以在 这里 查看生成的文档。如果文档构建失败，你可能需要 清除 test-docs 的构建环境。

对于本地测试，我能够在新的目录中通过重放我在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/google/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
目录