从源码构建 JAX：jaxlib、hermetic Python、测试与文档开发全指南

从源码构建 JAX：jaxlib、hermetic Python、测试与文档开发全指南

【免费下载链接】jaxComposable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/gh_mirrors/jax/jax

本文是 JAX（当前仓库）官方开发者文档 docs/developer.md 的深度展开版，系统讲解从源码构建 JAX 的完整流程：如何获取源码、构建或安装jaxlib（含 CUDA/ROCM/Windows/XLA 定制），如何利用 hermetic Python 机制锁定可复现的构建环境，如何通过 Bazel 或 pytest 运行测试，以及如何做类型检查、Lint 与文档维护。读完本文，你将能够在自己机器上搭建完整的 JAX 开发环境、构建多平台多后端的 wheel，并熟练运用仓库内提供的构建与测试工具链。

JAX 的构建模型：jaxlib 与 jax 的双包结构

JAX 的源码构建本质上分为两个独立步骤，这是理解整篇文档的前提：

1. 构建或安装jaxlib：这是 JAX 的 C++ 支持库，内含 XLA 编译器、运行时与各类内核实现，对应仓库中的 jaxlib/ 目录。
2. 安装jaxPython 包：即仓库根目录下的纯 Python 包，包含 jax/ 目录中的全部 Python 代码。

两者的安装是解耦的：如果你只修改 JAX 的 Python 部分，完全可以跳过 C++ 编译，直接用 pip 装一个预编译好的jaxlibwheel，再把jax以可编辑模式（pip install -e）装进环境，实现"改 Python 代码即时生效、C++ 部分用官方二进制"的开发体验。

获取源码的方式很直接（文档中约定以python作为 Python 3 解释器名称，部分系统需改用python3）：

git clone https://github.com/google/jax cd jax

本文中出现的所有build.py调用均指仓库根目录下build/build.py，它是整个构建流程的入口脚本（见下文源码剖析）。

构建或安装 jaxlib

方式一：用 pip 安装预编译 jaxlib（仅改 Python 代码时推荐）

如果开发工作只涉及 Python 层（例如 jax/_src/numpy/lax_numpy.py 这类纯 Python 实现），文档建议直接从 PyPI 安装预编译 wheel：

pip install jaxlib

GPU 与 TPU 支持所需的额外配置（CUDA wheel 选择、TPU 版本等）请参考 README.md 中关于 pip 安装的完整说明。

方式二：从源码构建 jaxlib

从源码构建需要先准备 C++ 编译工具链：

• Linux（Ubuntu/Debian）：

  sudo apt install g++ python python3-dev

• macOS：安装 Xcode 及 Xcode Command Line Tools。
• Windows：见下文专节。

一个关键设计是：构建过程不需要本地安装 Python 依赖——Bazel 会使用自己的 hermetic Python 解释器（详见"Managing hermetic Python"小节），你系统里的 Python 在构建期间会被忽略，只有build/build.py脚本本身由系统 Python 解释。

CPU 或 TPU 目标的标准构建命令：

python build/build.py pip install dist/*.whl # 安装 jaxlib（内含 XLA）

默认情况下 wheel 会输出到当前目录下的dist/子目录。要为与系统不同的 Python 版本构建 wheel，追加--python_version参数：

python build/build.py --python_version=3.12

build.py还支持大量配置选项，运行python build/build.py --help可查看全部参数，包括指定 CUDA/CuDNN 路径的方式。无论是否传--python_version，Bazel 侧始终使用 hermetic Python。

CUDA 支持：两种构建策略

文档给出了两条构建 CUDA 版 jaxlib 的路径：

1. 单 wheel 方案——把 CUDA 支持直接打进 jaxlib wheel：

   python build/build.py --enable_cuda

2. GPU 插件分离方案——生成三个 wheel（不含 CUDA 的 jaxlib、jax-cuda-plugin、jax-cuda-pjrt）：

   python build/build.py --enable_cuda --build_gpu_plugin --gpu_plugin_cuda_version=12

   gpu_plugin_cuda_version可设为 11 或 12。

从源码看，build/build.py在检测到--enable_cuda后会写入.jax_configure.bazelrc，其中包含build --config=cuda、TF_CUDA_PATHS、TF_CUDA_VERSION、TF_CUDNN_VERSION、TF_CUDA_COMPUTE_CAPABILITIES等 Bazel action env；--build_gpu_plugin则会追加build --config=cuda_plugin（build/build.py）。由此可见--enable_cuda本质上是把 CUDA 工具链通过 Bazel 配置注入到 XLA 的构建中，--build_gpu_plugin则切换到独立的 GPU 插件构建目标。

使用本地修改版 XLA 仓库构建

JAX 依赖 XLA，其源码位于独立的 XLA 仓库，默认使用 JAX 固定的（pinned）XLA 提交。当你在本地修改 XLA 时，有两种方式让构建使用你的副本：

1. Bazeloverride_repository特性，以命令行参数形式传给build.py：

   python build/build.py --bazel_options=--override_repository=xla=/path/to/xla

2. 直接修改 JAX 源码树根部的 WORKSPACE，将 XLA 指向不同的 tree。

若要回馈 XLA 改动，请向 XLA 仓库提交 PR。JAX 所固定的 XLA 版本会定期更新，尤其是每次jaxlib发布之前。

Windows 上的 jaxlib 构建要点

Windows 构建需要额外满足以下条件：

• C++ 工具链：安装 Visual Studio（2019 16.5 或更新版本）；若需 CUDA，按 NVIDIA CUDA 安装指南配置 CUDA 环境。
• 符号链接：JAX 构建使用符号链接，因此必须开启 Windows 的Developer Mode。
• Python 环境：可使用 Python 官方 Windows 安装包，或 Anaconda / Miniconda。
• MSYS2：Bazel 的部分目标依赖 bash 工具做脚本化，需安装 MSYS2，并安装patch与coreutils（后者提供realpath命令）：

  pacman -S patch coreutils

一切就绪后在 PowerShell 中确保bazel、patch、realpath可访问，激活 conda 环境后即可构建（以下示例开启了 CUDA，可按需调整路径与版本）：

python .\build\build.py ` --enable_cuda ` --cuda_path='C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v10.1' ` --cudnn_path='C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v10.1' ` --cuda_version='10.1' ` --cudnn_version='7.6.5'

如需带调试信息构建，追加--bazel_options='--copt=/Z7'。

为 AMD GPU 构建 ROCM 版 jaxlib

构建 ROCM 版需要若干 ROCM/HIP 库。以配置了 AMDapt仓库的 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 版本调整）：

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

AMD 的 XLA fork 可能包含上游没有的修复。若上游 XLA 遇到问题，可克隆 AMD fork 并覆盖构建所用的 XLA 仓库：

git clone https://github.com/ROCmSoftwarePlatform/xla.git python build/build.py --enable_rocm --rocm_path=/opt/rocm-5.7.0 \ --bazel_options=--override_repository=xla=/path/to/xla-rocm

与 CUDA 路径一致，--enable_rocm也会在.jax_configure.bazelrc中写入build --config=rocm与ROCM_PATH、TF_ROCM_AMDGPU_TARGETS等配置（build/build.py）。build/rocm/build_rocm.sh中还提供了针对 ROCM 环境的封装脚本可供参考。

管理 hermetic Python：可复现构建的核心机制

JAX 的所有 Bazel 构建与测试命令都依赖hermetic Python（基于 rules_python）：Bazel 自行管理 Python 解释器及其全部依赖，从而保证构建在 Linux、Windows、macOS 上行为一致，且与本地系统的 Python 完全隔离。

指定 Python 版本

运行build/build.py时，hermetic Python 的版本会自动匹配你用来执行脚本的 Python 版本；也可以用--python_version显式指定：

python build/build.py --python_version=3.12

底层由HERMETIC_PYTHON_VERSION环境变量控制。build/build.py会自动设置它；若直接运行 bazel，则需要手动设置，三种方式任选：

# 方式一：写入 .bazelrc build --repo_env=HERMETIC_PYTHON_VERSION=3.12 # 方式二：直接在构建命令中传递 bazel build <target> --repo_env=HERMETIC_PYTHON_VERSION=3.12 # 方式三：在 shell 中全局导出 export HERMETIC_PYTHON_VERSION=3.12

值得注意的是，build.py本身也会把HERMETIC_PYTHON_VERSION写入生成的.jax_configure.bazelrc（build/build.py），这意味着用build.py构建时版本一致性是自动保证的，无需手动干预。

因为不同 Python 版本共享与解释器无关的构建缓存，你可以在同一台机器上通过切换--python_version依次对不同版本的 Python 构建和测试，之前的缓存会被保留复用。

指定 Python 依赖：requirements 锁定文件

为保证构建可复现，Bazel 构建期间 JAX 的全部 Python 依赖都被钉死到特定版本。完整的依赖传递闭包及其哈希记录在build/requirements_lock_<python version>.txt文件中（例如 Python 3.12 对应 build/requirements_lock_3_12.txt；当前仓库提供了 3.10～3.13 的锁定文件）。

直接依赖清单维护在 build/requirements.in 中，其中既有运行时依赖（numpy、scipy、ml_dtypes、opt_einsum、zstandard、etils[epath]），也有测试依赖（通过-r test-requirements.txt引入）。更新锁定文件的命令：

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

该命令底层调用pip-compile（pip-tools）。如果希望有更多控制，也可以直接运行等价的 Bazel 命令：

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

其中3.12是你希望更新的 Python 版本。由于底层仍是pip与pip-compile，这两个工具支持的绝大多数命令行参数都会被接受，例如想让更新器考虑预发布版本：

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

依赖本地 wheel

如果需要依赖本地.whl文件（比如你刚构建好的 jaxlib wheel），把 wheel 的路径追加到 build/requirements.in 并重新运行更新器即可：

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

依赖 nightly wheel

要针对最新、可能不稳定的一组依赖构建测试，使用 nightly 版本更新器：

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

与常规更新器的区别在于：它默认接受预发布、dev 与 nightly 包；会额外搜索pypi.anaconda.org/scientific-python-nightly-wheels作为 extra index url；且最终锁定文件不写入哈希。

用预发布版 Python 构建（进阶）

JAX 开箱即支持所有当前正式发布的 Python 版本；若要针对尚未正式发布的版本构建测试，按以下步骤：

1. 安装编译 Python 解释器及其关键包（如 numpy/scipy）所需的系统包。典型 Debian 系统：

   sudo apt-get update sudo apt-get build-dep python3 -y sudo apt-get install pkg-config zlib1g-dev libssl-dev -y # 构建 scipy 需要： sudo apt-get install libopenblas-dev -y

2. 检查 WORKSPACE，确认其中存在指向目标 Python 版本的custom_python_interpreter()条目。
3. 构建 Python 解释器：

   bazel build @python_dev//:python_dev

   默认用 GCC 编译；若想用 clang，可设置对应环境变量（例如--repo_env=CC=/usr/lib/llvm-17/bin/clang --repo_env=CXX=/usr/lib/llvm-17/bin/clang++）。

4. 把生成的python_register_toolchains()片段写入 WORKSPACE。上一步命令的末尾会打印一段代码片段，将其复制到python_init_toolchains()条目之后（新增版本）或替换之（如用自建 3.12 替换默认 3.12）。片段已按你的实际环境生成，可直接使用，也可自定义（比如把 Python 的.tgz改为远程下载地址）。
5. 确保 WORKSPACE 的python_init_repositories()的requirements参数包含该版本条目，例如 Python 3.13 应有"3.13": "//build:requirements_lock_3_13.txt"。
6. （推荐）对不稳定版本预构建全部 Python 依赖：

   bazel build //build:all_py_deps --repo_env=HERMETIC_PYTHON_VERSION="3.13"

   这会让 pip 对尚无二进制分发的包（如 numpy、scipy、matplotlib、zstandard）从源码拉取构建。建议在实际 JAX 构建之前独立执行此步骤，以免两者互相冲突。例如 JAX 通常用 clang 构建，而matplotlib从源码构建时假设使用 GCC，clang 会因 LTO（-flto）行为差异导致构建失败。若针对稳定版本 Python 构建、或所有依赖都有现成二进制分发，可跳过此步。

7. 构建完成后，正常执行构建/测试命令即可，只需确保HERMETIC_PYTHON_VERSION指向你的新版本。
8. 关于锁定文件更新的提醒：对预发布版本 Python 直接更新requirements_lock_<python_version>.txt很可能失败——仓库里没有匹配的二进制包时，pip-compile会尝试从源码构建，这比pip安装更严格，更容易失败。推荐做法是：为最新稳定版本（如 3.12）生成不带哈希的锁定文件，再复制给不稳定版本（如 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" # 根据依赖对新版本 Python 的兼容程度，你可能需要手动编辑最终的锁定文件

安装 jax Python 包

jaxlib就绪后，在仓库根部用可编辑模式安装jax：

pip install -e . # 安装 jax

要从 GitHub 升级到最新版，只需在仓库根部执行git pull，然后按需重新运行build.py或升级jaxlib。通常不需要重装jax——pip install -e已经在 site-packages 和仓库之间建立了符号链接。

运行测试

JAX 测试支持 Bazel 与 pytest 两种机制。

使用 Bazel 运行测试

首先配置构建：

python build/build.py --configure_only

可按需向build.py追加其他配置选项。默认情况下 Bazel 测试使用从源码构建的 jaxlib，运行方式：

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

//tests:gpu_tests与//tests:tpu_tests目标同样可用（需要相应硬件）。

如果希望改用预装的 jaxlib（而非每次重新构建），需要先把 jaxlib 装进 hermetic Python。安装指定版本（以jaxlib >= 0.4.26为例）：

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

或从本地 wheel 安装（假设 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

hermetic 环境装好 jaxlib 后，用如下命令跑测试：

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

多加速器测试：部分测试面向多 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

单加速器测试并行加速：可在多个加速器上并行运行，每个加速器同时跑多个测试。以 2 块 GPU、每块 4 个并发任务为例：

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

这里的 build/parallel_accelerator_execute.sh 就是仓库中提供的多加速器执行包装脚本，通过JAX_ACCELERATOR_COUNT与JAX_TESTS_PER_ACCELERATOR两个环境变量控制资源划分。

使用 pytest 运行测试

先安装依赖：

pip install -r build/test-requirements.txt

build/test-requirements.txt 中已包含pytest-xdist，因此可以从仓库根部直接并行运行全部测试：

pytest -n auto tests

-n auto由 pytest-xdist 提供，会按 CPU 核数自动并行。

控制测试行为

JAX 会组合式地生成测试用例，每个测试默认生成并校验 10 个用例，可用JAX_NUM_GENERATED_CASES环境变量调整（CI 自动化测试默认使用 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

• 运行单个测试文件可看到更详细的用例信息：

  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 notebook 无报错。

Doctests

JAX 用 pytest 的 doctest 模式测试文档中的代码示例：

pytest docs

此外还以doctest-modules模式确保函数 docstring 中的示例可运行，例如：

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

注意：对整个包执行 doctest 时，有若干文件被标记为跳过，细节可在 CI 工作流配置（仓库中的.github/workflows/ci-build.yaml）中查看。

类型检查

JAX 使用mypy检查类型标注，与 CI 完全一致的本机检查命令：

pip install mypy mypy --config=pyproject.toml --show-error-codes jax

mypy 配置位于仓库根部的 pyproject.toml。也可以借助 .pre-commit-config.yaml 中定义的 pre-commit 钩子，对 git 暂存区文件自动运行与 GitHub CI 相同版本的 mypy：

pre-commit run mypy

Linting

JAX 使用ruff保证代码质量，本机检查：

pip install ruff ruff jax

同样支持通过 pre-commit 对暂存文件自动执行（与 CI 相同版本）：

pre-commit run ruff

更新文档

重建文档（Sphinx）

安装文档构建依赖：

pip install -r docs/requirements.txt

然后构建 HTML 文档：

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

由于构建会执行文档源中的大量 notebook，耗时可能很长；跳过 notebook 执行可加速：

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

生成结果位于docs/build/html/index.html。-j auto控制构建并行度，也可换成具体数字以指定 CPU 核心数。

维护 notebook（ipynb 与 md 双格式）

文档使用 jupytext 中维护两份同步的 notebook：ipynb格式可直接在 Colab 打开执行，md格式则便于版本控制中查看 diff（仓库中如 docs/notebooks/thinking_in_jax.md 与同名.ipynb即为一对）。

• 编辑 ipynb：改动较大（涉及代码与输出）时，建议在 Jupyter 或 Colab 中编辑，完成后Run all cells再Download ipynb，并按上文sphinx-build方式验证可执行。
• 编辑 md：仅修改文本内容时直接编辑.md版本更方便。
• 同步双版本：编辑完任一版本后，用 jupytext 同步（版本应与 .pre-commit-config.yaml 中指定的一致）：

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

  校验同步是否正确的 pre-commit 钩子：

  git add docs -u # pre-commit 只处理暂存区文件 pre-commit run jupytext

• 新建 notebook：若要纳入 jupytext 同步，先设置格式（会在 notebook 中写入jupytext元数据字段，jupytext --sync据此识别）：

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

• Sphinx 构建中的 notebook：部分 notebook 会在预提交检查与 Read the Docs 构建中自动执行，单元格报错会导致构建失败。如果报错是有意为之，可在.ipynb中手动为单元格打上raises-exceptions元数据（保存时会被保留）。含长计算等场景的 notebook 会通过 docs/conf.py 的exclude_patterns排除出构建。

Read the Docs 上的文档构建

JAX 的自动化文档位于 jax.readthedocs.io，构建由项目级 Read the Docs 设置驱动：代码推送到main分支即触发文档构建。每次构建由 .readthedocs.yml 与 docs/conf.py 驱动——前者定义了构建镜像（ubuntu-22.04）、Python 版本（3.10）、Sphinx 配置路径（docs/conf.py，并开启fail_on_warning: true）以及额外格式（htmlzip）和docs/requirements.txt依赖安装。

• 推送到test-docs分支也会自动构建，可用于预演文档效果。
• 本地复现 Read the Docs 构建：可参考其构建日志在全新目录中依次执行依赖安装与sphinx-build命令（例如创建虚拟环境、git clone --no-single-branch --depth 50拉取仓库、切换test-docs分支、安装docs/requirements.txt，最后运行python \which sphinx-build` -T -E -b html -d _build/doctrees-readthedocs -D language=en . _build/html`）。

附：build.py 底层实现速览

整个构建体系的入口是 build/build.py（共 749 行），从源码可以确认几个关键实现细节：

• Bazel 自动获取与校验：脚本内置了各平台（Linux x86_64/aarch64、Darwin x86_64/arm64、Windows AMD64）的 Bazel 6.5.0 下载映射，找不到本地 bazel 时自动下载，并用内置 SHA256 校验二进制完整性（build/build.py）。
• Bazel 版本下限：get_bazel_path要求 bazel 版本 >= 6.5.0，否则报错退出（build/build.py）。
• Python 版本下限：check_python_version要求 Python 3.10 或更新版本（build/build.py），这与仓库提供requirements_lock_3_10.txt~requirements_lock_3_13.txt的覆盖范围一致。
• 配置生成：write_bazelrc根据--enable_cuda、--enable_rocm、--build_gpu_plugin、--use_clang、--target_cpu_features、--python_version等参数生成.jax_configure.bazelrc，把 CUDA/ROCM 路径、编译配置与HERMETIC_PYTHON_VERSION注入 Bazel（build/build.py）。
• CPU 特性选项：--target_cpu_features支持release（x86-64 下启用 AVX）、native（-march=native，Windows 不支持）、default（交给编译器默认行为）（build/build.py）。

理解这些底层机制，有助于你在自定义构建参数或排查构建问题时定位到build.py的具体分支逻辑。

【免费下载链接】jaxComposable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/gh_mirrors/jax/jax