从源代码构建

从源代码构建#

首先，获取 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 wheel，或 (2) 使用 python build/build.py --enable_cuda --build_gpu_plugin --gpu_plugin_cuda_version=12 生成三个 wheel（没有 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 Redistributables，您可以使用以下命令

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

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

    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"

    有关完整说明列表，请参见 XLA 文档。

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

使用修改后的 XLA 存储库从源代码构建 jaxlib#

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）上行为一致，并且与本地系统的具体情况隔离开来，我们依赖于封闭的 Python（由 rules_python 提供，请参见 工具链注册 获取详细信息）用于通过 Bazel 执行的所有构建和测试命令。这意味着您的系统 Python 安装将在构建过程中被忽略，Python 解释器本身以及所有 Python 依赖项将由 bazel 直接管理。

指定 Python 版本#

当您运行 build/build.py 工具时，封闭 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，并且不会在生成的 requirements 锁定文件中添加哈希值。

使用预发布的 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 中指定的任何版本定义）。为了使构建成功，请确保您选择执行构建的隔离 Python 已经存在于您的配置中（实际版本无关紧要，只要它是一个有效的版本）。默认情况下，Python 二进制文件将使用 GCC 编译器构建。如果您希望使用 clang 构建它，您需要设置相应的环境变量来执行此操作（例如 --repo_env=CC=/usr/lib/llvm-17/bin/clang --repo_env=CXX=/usr/lib/llvm-17/bin/clang++）。

  3. 检查上一个命令的输出。在它的最后，您将找到一个用于 python_register_toolchains() 条目的代码片段，其中包含您新构建的 Python。将该代码片段复制到您的 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 构建）以获取所有不稳定的 Python 版本，以避免在构建 JAX 本身和构建其 Python 依赖项之间发生冲突。例如，我们通常使用 clang 构建 JAX，但是从源代码使用 clang 构建 matplotlib 会开箱即用地失败，这是因为 GCC 和 clang 之间的 LTO（链接时优化，由 -flto 标志触发）行为存在差异，并且 matplotlib 默认情况下假定使用 GCC。如果您针对稳定的 Python 版本进行构建，或者通常您不希望任何 Python 依赖项从源代码构建（即对应 Python 版本的二进制发行版已经存在于存储库中），则不需要执行此步骤。

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

  7. 注意，如果您正在构建预发布版本的 Python，则使用您新构建的 Python 更新 requirements_lock_<python_version>.txt 文件可能会失败，因为包存储库中将没有匹配的二进制包。如果没有可用的二进制包，pip-compile 将继续从源代码构建它们，这可能会失败，因为它的限制比在 pip 安装期间执行相同操作更严格。更新不稳定 Python 版本的 requirements 锁定文件的推荐方法是：更新最新稳定版本的 requirements（例如 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

或者，要从本地轮子安装 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

可以使用环境变量控制许多测试行为（见下文）。可以使用 --test_env=FLAG=value 标志将环境变量传递给 Bazel 的 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。通常情况下，假设将使用多个示例输入进行测试，如果测试失败，它将尝试找到一个仍然会导致失败的更小的示例：查看测试失败以查找类似于下面的行，并在消息中添加提到的装饰器

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 来测试文档中的代码示例。你可以使用以下命令运行它

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

Linting#

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

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

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 并从本地仓库 Upload。根据需要更新它，Run all cells，然后 Download ipynb。你可能需要使用上面解释的 sphinx-build 测试它是否正常执行。

编辑 md#

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

同步笔记本#

在编辑笔记本的 ipynb 或 md 版本后，你可以使用 jupytext 通过对更新后的笔记本运行 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

创建新的笔记本#

如果你要将新的笔记本添加到文档中，并且希望使用此处讨论的 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.ac.cn/。

文档构建由 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
目录