CI: 12/02/24 upstream sync

.bazelrc

@@ -183,6 +183,7 @@ build:macos_cache_push --config=macos_cache --remote_upload_local_results=true -
 build:ci_linux_x86_64 --config=avx_linux --config=avx_posix
 build:ci_linux_x86_64 --config=mkl_open_source_only
 build:ci_linux_x86_64 --config=clang --verbose_failures=true
+build:ci_linux_x86_64 --color=yes
 
 # TODO(b/356695103): We do not have a CPU only toolchain so we use the CUDA
 # toolchain for both CPU and GPU builds.
@@ -203,6 +204,7 @@ build:ci_linux_x86_64_cuda --config=ci_linux_x86_64
 # Linux Aarch64 CI configs
 build:ci_linux_aarch64_base --config=clang --verbose_failures=true
 build:ci_linux_aarch64_base --action_env=TF_SYSROOT="/dt10"
+build:ci_linux_aarch64_base --color=yes
 
 build:ci_linux_aarch64 --config=ci_linux_aarch64_base
 build:ci_linux_aarch64 --host_crosstool_top="@ml2014_clang_aarch64_config_aarch64//crosstool:toolchain"
@@ -221,18 +223,21 @@ build:ci_linux_aarch64_cuda --action_env=CLANG_CUDA_COMPILER_PATH="/usr/lib/llvm
 build:ci_darwin_x86_64 --macos_minimum_os=10.14
 build:ci_darwin_x86_64 --config=macos_cache_push
 build:ci_darwin_x86_64 --verbose_failures=true
+build:ci_darwin_x86_64 --color=yes
 
 # Mac Arm64 CI configs
 build:ci_darwin_arm64 --macos_minimum_os=11.0
 build:ci_darwin_arm64 --config=macos_cache_push
 build:ci_darwin_arm64 --verbose_failures=true
+build:ci_darwin_arm64 --color=yes
 
 # Windows x86 CI configs
 build:ci_windows_amd64 --config=avx_windows
 build:ci_windows_amd64 --compiler=clang-cl --config=clang --verbose_failures=true
 build:ci_windows_amd64 --crosstool_top="@xla//tools/toolchains/win/20240424:toolchain"
 build:ci_windows_amd64 --extra_toolchains="@xla//tools/toolchains/win/20240424:cc-toolchain-x64_windows-clang-cl"
 build:ci_windows_amd64 --host_linkopt=/FORCE:MULTIPLE --linkopt=/FORCE:MULTIPLE
+build:ci_windows_amd64 --color=yes
 
 # #############################################################################
 # RBE config options below. These inherit the CI configs above and set the

.github/workflows/asan.yaml

@@ -25,14 +25,8 @@ jobs:
       run:
         shell: bash -l {0}
     steps:
-      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
-        with:
-          path: jax
-      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
-        with:
-          repository: python/cpython
-          path: cpython
-          ref: v3.13.0
+      # Install git before actions/checkout as otherwise it will download the code with the GitHub
+      # REST API and therefore any subsequent git commands will fail.
       - name: Install clang 18
         env:
           DEBIAN_FRONTEND: noninteractive
@@ -42,6 +36,14 @@ jobs:
             zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev curl git \
             libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev \
             libffi-dev liblzma-dev
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        with:
+          path: jax
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        with:
+          repository: python/cpython
+          path: cpython
+          ref: v3.13.0
       - name: Build CPython with ASAN enabled
         env:
           ASAN_OPTIONS: detect_leaks=0
@@ -65,7 +67,7 @@ jobs:
         run: |
           source ${GITHUB_WORKSPACE}/venv/bin/activate
           cd jax
-          python build/build.py \
+          python build/build.py build --wheels=jaxlib --verbose \
             --bazel_options=--color=yes \
             --bazel_options=--copt=-fsanitize=address \
             --clang_path=/usr/bin/clang-18

.github/workflows/ci-build.yaml

@@ -101,7 +101,7 @@ jobs:
 
   documentation:
     name: Documentation - test code snippets
-    runs-on: ubuntu-latest
+    runs-on: ROCM-Ubuntu
     timeout-minutes: 10
     strategy:
       matrix:
@@ -164,8 +164,7 @@ jobs:
         pip install -r docs/requirements.txt
     - name: Render documentation
       run: |
-        sphinx-build --color -W --keep-going -b html -D nb_execution_mode=off docs docs/build/html
-
+        sphinx-build -j auto --color -W --keep-going -b html -D nb_execution_mode=off docs docs/build/html
 
   jax2tf_test:
     name: "jax2tf_test (py ${{ matrix.python-version }} on ${{ matrix.os }}, x64=${{ matrix.enable-x64}})"

.github/workflows/cloud-tpu-ci-nightly.yml

@@ -24,10 +24,10 @@ jobs:
     strategy:
       fail-fast: false # don't cancel all jobs on failure
       matrix:
-        jaxlib-version: ["pypi_latest", "nightly", "nightly+oldest_supported_libtpu"]
+        jaxlib-version: ["head", "pypi_latest", "nightly", "nightly+oldest_supported_libtpu"]
         tpu: [
-          # {type: "v3-8", cores: "4"}, # Enable when we have the v3/v4 type available
-          # {type: "v4-8", cores: "4", runner: "linux-x86-ct4p-240-4tpu"},
+          # {type: "v3-8", cores: "4"}, # Enable when we have the v3 type available
+          {type: "v4-8", cores: "4", runner: "linux-x86-ct4p-240-4tpu"},
           {type: "v5e-8", cores: "8", runner: "linux-x86-ct5lp-224-8tpu"}
         ]
         python-version: ["3.10"]
@@ -47,14 +47,34 @@ jobs:
       # mandates using a specific commit for non-Google actions. We use
       # https://github.com/sethvargo/ratchet to pin specific versions.
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+      # Checkout XLA at head, if we're building jaxlib at head.
+      - name: Checkout XLA at head
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+        if: ${{ matrix.jaxlib-version == 'head' }}
+        with:
+          repository: openxla/xla
+          path: xla
       - name: Install JAX test requirements
         run: |
           $PYTHON -m pip install -U -r build/test-requirements.txt
           $PYTHON -m pip install -U -r build/collect-profile-requirements.txt
       - name: Install JAX
         run: |
           $PYTHON -m pip uninstall -y jax jaxlib libtpu
-          if [ "${{ matrix.jaxlib-version }}" == "pypi_latest" ]; then
+          if [ "${{ matrix.jaxlib-version }}" == "head" ]; then
+            # Build and install jaxlib at head
+            $PYTHON build/build.py --bazel_options=--config=rbe_linux_x86_64 \
+                    --bazel_options="--override_repository=xla=$(pwd)/xla" \
+                    --bazel_options=--color=yes
+            $PYTHON -m pip install dist/*.whl
+
+            # Install "jax" at head
+            $PYTHON  -m pip install -U -e .
+
+            # Install libtpu
+            $PYTHON -m pip install --pre libtpu \
+              -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+          elif [ "${{ matrix.jaxlib-version }}" == "pypi_latest" ]; then
             $PYTHON -m pip install .[tpu] \
               -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
 

.github/workflows/wheel_win_x64.yml

@@ -40,7 +40,7 @@ jobs:
           python -m pip install -r build/test-requirements.txt
           python -m pip install --upgrade numpy==2.0.0 scipy==1.13.1
           "C:\\msys64\\;C:\\msys64\\usr\\bin\\;" >> $env:GITHUB_PATH
-          python.exe build\build.py `
+          python.exe build\build.py build --wheels=jaxlib `
             --bazel_options=--color=yes `
             --bazel_options=--config=win_clang `
             --verbose

.github/workflows/windows_ci.yml

@@ -49,9 +49,10 @@ jobs:
           python -m pip install -r build/test-requirements.txt
           python -m pip install --upgrade numpy==2.0.0 scipy==1.13.1
           "C:\\msys64\\;C:\\msys64\\usr\\bin\\;" >> $env:GITHUB_PATH
-          python.exe build\build.py `
+          python.exe build\build.py build --wheels=jaxlib `
             --bazel_options=--color=yes `
-            --bazel_options=--config=win_clang
+            --bazel_options=--config=win_clang `
+            --verbose
 
       - uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882  # v4.4.3
         with:

CHANGELOG.md

@@ -55,6 +55,11 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
         `platforms` instead.
   * Hashing of tracers, which has been deprecated since version 0.4.30, now
     results in a `TypeError`.
+  * Refactor: JAX build CLI (build/build.py) now uses a subcommand structure and
+    replaces previous build.py usage. Run `python build/build.py --help` for
+    more details. Brief overview of the new subcommand options:
+    * `build`: Builds JAX wheel packages. For e.g., `python build/build.py build --wheels=jaxlib,jax-cuda-pjrt`
+    * `requirements_update`: Updates requirements_lock.txt files.
   * {func}`jax.scipy.linalg.toeplitz` now does implicit batching on multi-dimensional
     inputs. To recover the previous behavior, you can call {func}`jax.numpy.ravel`
     on the function inputs.