從原始碼建置#

首先，取得 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

請參閱 JAX readme 以取得有關 pip 安裝的完整指南（例如，針對 GPU 和 TPU 支援）。

從原始碼建置 jaxlib#

警告

雖然通常應該可以使用大多數現代編譯器從原始碼編譯 jaxlib，但建置僅使用 clang 進行測試。我們歡迎提交 Pull Request 以改進對不同工具鏈的支援，但其他編譯器未獲得積極支援。

若要從原始碼建置 jaxlib，您也必須安裝一些先決條件

  • C++ 編譯器

    如上方方框中所述，最好使用最新版本的 clang（在撰寫本文時，我們測試的版本為 18），但其他編譯器（例如 g++ 或 MSVC）也可能適用。

    在 Ubuntu 或 Debian 上，您可以按照 LLVM 文件中的指示安裝最新穩定版本的 clang。

    如果您在 Mac 上建置，請確保已安裝 XCode 和 XCode 命令列工具。

    請參閱下方以取得 Windows 建置指示。

  • Python：用於執行建置輔助程式碼。請注意，無需在本機安裝 Python 依賴項，因為您的系統 Python 將在建置期間被忽略；請查看 管理隔離的 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 建置都將始終使用隔離的 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 始終會使用其自己的隔離 Python 解譯器和依賴項，只有 build/build.py 腳本本身將由您的系統 Python 解譯器處理。依預設，wheel 會寫入目前目錄的 dist/ 子目錄。

  • 從 v.0.4.32 開始的 JAX 版本：您可以在組態選項中提供自訂 CUDA 和 CUDNN 版本。Bazel 將下載它們並將其用作目標依賴項。

    若要下載特定版本的 CUDA/CUDNN 重新發行套件，您可以使用 --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 重新發行版本。

    若要指向本機檔案系統上的 CUDA/CUDNN/NCCL 重新發行套件，您可以使用以下命令

    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 文件中的完整指示列表。

  • 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 build --wheels=jaxlib --local_xla_path=/path/to/xla

  • 修改 JAX 原始碼樹狀結構根目錄中的 WORKSPACE 檔案，以指向不同的 XLA 樹狀結構。

若要將變更貢獻回 XLA，請將 PR 傳送至 XLA 儲存庫。

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

管理隔離的 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 build --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 執行建置和測試，只需在執行之間切換 --python_version 的值即可。先前建置中的所有與 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

指定本機 wheel 的依賴項#

依預設，建置會掃描儲存庫根目錄中的 dist 目錄，以尋找要包含在依賴項清單中的任何本機 .whl 檔案。如果 wheel 是 Python 版本特定的，則只會包含與選取的 Python 版本相符的 wheel。

整體本機 wheel 搜尋和選取邏輯由 python_init_repositories() 巨集（直接從 WORKSPACE 檔案呼叫）的引數控制。您可以使用 local_wheel_dist_folder 來變更包含本機 wheel 的資料夾位置。使用 local_wheel_inclusion_list 和 local_wheel_exclusion_list 引數來指定應包含和/或排除在搜尋中的 wheel（它支援基本萬用字元比對）。

如有必要，您也可以手動依賴本機 .whl 檔案，繞過自動本機 wheel 搜尋機制。例如，若要依賴您新建置的 jaxlib wheel，您可以將 wheel 的路徑新增至 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

指定 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，並且不會將雜湊值放入產生的依賴項鎖定檔案中。

自訂隔離的 Python (進階用法)#

我們開箱即用地支援所有目前的 Python 版本，因此，除非您的工作流程有非常特殊的需求（例如使用您自己的自訂 Python 解譯器的能力），否則您可以安全地完全略過本節。

簡而言之，如果您依賴非標準的 Python 工作流程，您仍然可以在隔離的 Python 設定中實現高度的靈活性。從概念上講，與非隔離的情況相比，只會有一個差異：您需要以檔案而不是安裝的角度來思考（即思考您的建置實際依賴哪些檔案，而不是需要在您的系統上安裝哪些檔案），其餘的幾乎相同。

因此，實際上，若要完全控制您的 Python 環境（無論是否隔離），您需要能夠執行以下三件事

  1. 指定要使用的 python 解譯器（即選取實際的 python 或 python3 二進位檔案以及與其位於同一資料夾中的程式庫）。

  2. 指定 Python 依賴項清單（例如 numpy）及其確切版本。

  3. 能夠輕鬆地在清單中新增/移除/更新依賴項。每個依賴項本身也可以是自訂的（例如自行建置）。

您已經知道如何在非隔離的 Python 環境中執行上述所有步驟，以下說明如何在隔離的環境中執行相同的步驟（透過以檔案而不是安裝的角度來處理它）

  1. 不要安裝 Python，而是以 tar 或 zip 檔案取得 Python 解譯器。根據您的情況，您可以簡單地提取許多現有的解譯器之一（例如 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 tarbal 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（請注意，鎖定檔案必須是隔離的，但如果您願意，您可以隨時從非隔離的 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，則它會完全控制您的依賴項列表，並且指定本機 wheels 的依賴項中描述的自動本機 wheels 解析邏輯將被停用，以避免干擾它。

就是這樣。總結一下：如果您有一個包含 Python 直譯器的封存檔，以及一個包含您的依賴項完整傳遞閉包的 requirements_lock.txt 檔案，那麼您就可以完全控制您的 Python 環境。

自訂 hermetic Python 範例#

請注意，對於以下所有範例，您也可以全域設定環境變數（即在您的 shell 中使用 export 而不是命令的 --repo_env 參數），以便透過 build/build.py 呼叫 bazel 可以正常運作。

從網際網路使用自訂 Python 3.13 進行建置，使用已在此 repo 中簽入的預設 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（假設鎖定檔在執行命令之前已放入此 repo 的 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"

請注意，您可以有多個不同的 requirement_lock.txt 檔案，對應於相同的 Python 版本，以支援不同的情境。您可以透過指定 HERMETIC_PYTHON_VERSION 來控制選擇哪一個。例如，在 WORKSPACE 檔案中

requirements = {
  "3.10": "//build:requirements_lock_3_10.txt",
  "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 dependendencies:
bazel test <target> --repo_env=HERMETIC_PYTHON_VERSION=3.13-scenario1

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

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

# To build with scenario1 dependendencies 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 以用於 GPU

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 也可用。

若要使用預先安裝的 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-manylinux2014_x86_64.whl)" >> build/requirements.in
python build/build.py requirements_update --python_version=3.12

一旦您 hermetically 安裝了 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。通常，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 中找到執行 doctest 的最新命令。例如，您可以執行

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，以確保函式 docstring 中的程式碼範例可以正確執行。您可以使用例如以下命令在本機執行此操作

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

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 --sync，使用 jupytext 來同步這兩個版本；例如

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" metadata 欄位新增到筆記本檔案來運作，該欄位指定所需的格式，以及在調用 jupytext --sync 命令時識別的格式。

Sphinx 建置中的筆記本#

某些筆記本會自動建置為預先提交檢查和 Read the docs 建置的一部分。如果儲存格引發錯誤，建置將會失敗。如果錯誤是故意的，您可以捕獲它們，或使用 raises-exceptions metadata 標記儲存格（範例 PR）。您必須在 .ipynb 檔案中手動新增此 metadata。當其他人重新儲存筆記本時，它將被保留。

我們從建置中排除了一些筆記本，例如，因為它們包含長時間的計算。請參閱 conf.py 中的 exclude_patterns。

在 readthedocs.io 上建置文件#

JAX 的自動產生文件位於 https://jax.dev.org.tw/。

文件建置由 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/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