Merge branch 'main' into srnitin/test-scripts-workflows

Author: nitins17

.bazelrc

@@ -1,6 +1,12 @@
 # #############################################################################
 # All default build options below. These apply to all build commands.
 # #############################################################################
+# TODO: Enable Bzlmod
+common --noenable_bzlmod
+
+# TODO: Migrate for https://github.com/bazelbuild/bazel/issues/7260
+common --noincompatible_enable_cc_toolchain_resolution
+
 # Make Bazel print out all options from rc files.
 common --announce_rc
 
@@ -41,6 +47,13 @@ build:linux --copt=-Wno-array-parameter
 build:macos --config=posix
 build:macos --apple_platform_type=macos
 
+# Bazel 7.0.0 no longer supports dynamic symbol lookup on macOS. To resolve
+# undefined symbol errors in macOS arm64 builds, explicitly add the necessary
+# linker flags until dependencies are well defined. See
+# https://github.com/bazelbuild/bazel/issues/19730.
+build:macos --linkopt=-Wl,-undefined,dynamic_lookup
+build:macos --host_linkopt=-Wl,-undefined,dynamic_lookup
+
 # Windows has a relatively short command line limit, which JAX has begun to hit.
 # See https://docs.bazel.build/versions/main/windows.html
 build:windows --features=compiler_param_file
@@ -270,6 +283,12 @@ build:resultstore --bes_instance_name="tensorflow-testing"
 build:resultstore --bes_results_url="https://source.cloud.google.com/results/invocations"
 build:resultstore --bes_timeout=600s
 
+# Configs for RBE cache. When using resultstore, we need to use these configs
+# as well to ensure that the logs that get uploaded to resultstore can be read
+# without any errors.
+build:rbe_cache --remote_cache=remotebuildexecution.googleapis.com
+build:rbe_cache --remote_instance_name=projects/tensorflow-testing/instances/default_instance
+
 build:rbe --config=resultstore
 build:rbe --repo_env=BAZEL_DO_NOT_DETECT_CPP_TOOLCHAIN=1
 build:rbe --define=EXECUTOR=remote

.bazelversion

@@ -1 +1 @@
-6.5.0
+7.4.1

.github/workflows/asan.yaml

@@ -17,7 +17,7 @@ on:
 jobs:
   asan:
     # Don't execute in fork due to runner type
-    if: github.repository == 'jax-ml/jax'
+    if: ${{ github.repository == 'jax-ml/jax' }}
     runs-on: linux-x86-n2-64
     container:
       image: index.docker.io/library/ubuntu@sha256:b359f1067efa76f37863778f7b6d0e8d911e3ee8efa807ad01fbf5dc1ef9006b # ratchet:ubuntu:24.04
@@ -62,7 +62,8 @@ jobs:
         run: |
           source ${GITHUB_WORKSPACE}/venv/bin/activate
           cd jax
-          pip install -r build/test-requirements.txt
+          pip install uv~=0.5.30
+          uv pip install -r build/test-requirements.txt
       - name: Build and install JAX
         env:
           ASAN_OPTIONS: detect_leaks=0
@@ -73,8 +74,8 @@ jobs:
             --bazel_options=--color=yes \
             --bazel_options=--copt=-fsanitize=address \
             --clang_path=/usr/bin/clang-18
-          pip install dist/jaxlib-*.whl
-          pip install -e .
+          uv pip install dist/jaxlib-*.whl \
+            -e .
       - name: Run tests
         env:
           ASAN_OPTIONS: detect_leaks=0
@@ -89,4 +90,4 @@ jobs:
           echo "JAX_ENABLE_X64=$JAX_ENABLE_X64"
           echo "JAX_SKIP_SLOW_TESTS=$JAX_SKIP_SLOW_TESTS"
           # The LD_PRELOAD works around https://github.com/google/sanitizers/issues/934#issuecomment-649516500
-          LD_PRELOAD=/lib/x86_64-linux-gnu/libstdc++.so.6 python -m pytest -n auto --tb=short --maxfail=20 tests
+          LD_PRELOAD=/lib/x86_64-linux-gnu/libstdc++.so.6 python -m pytest -n 32 --tb=short --maxfail=20 tests

.github/workflows/bazel_cuda_non_rbe.yml

@@ -47,7 +47,7 @@ jobs:
         # Explicitly set the shell to bash
         shell: bash
     runs-on: ${{ inputs.runner }}
-    container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-cudnn9.1-cuda12.3:720686788"
+    container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-cuda12.3-cudnn9.1:latest"
 
     env:
       JAXCI_HERMETIC_PYTHON_VERSION: ${{ inputs.python }}

.github/workflows/build_artifacts.yml

@@ -17,7 +17,7 @@ on:
         options:
         - "linux-x86-n2-16"
         - "linux-arm64-c4a-64"
-        - "windows-x86-n2-16"
+        - "windows-x86-n2-64"
       artifact:
         description: "Which JAX artifact to build?"
         type: choice
@@ -125,19 +125,19 @@ jobs:
         with:
           repository: openxla/xla
           path: jax/xla
-      - name: Enable RBE if building on Linux x86 or Windows x86
-        if: contains(inputs.runner, 'linux-x86') || contains(inputs.runner, 'windows-x86')
+      - name: Enable RBE if building on Linux x86
+        if: contains(inputs.runner, 'linux-x86')
         run: echo "JAXCI_BUILD_ARTIFACT_WITH_RBE=1" >> $GITHUB_ENV
-      - name: Enable Bazel remote cache (with writes enabled) if building on Linux Aarch64
-        if: contains(inputs.runner, 'linux-arm64')
+      - name: Enable Bazel remote cache (with writes enabled) if building on Linux Aarch64 or Windows x86
+        if: contains(inputs.runner, 'linux-arm64') || contains(inputs.runner, 'windows-x86')
         run: echo "JAXCI_WRITE_TO_BAZEL_REMOTE_CACHE=1" >> $GITHUB_ENV
       # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main
         with:
           halt-dispatch-input: ${{ inputs.halt-for-connection }}
       - name: Build ${{ inputs.artifact }}
-        timeout-minutes: 30
+        timeout-minutes: 60
         run: ./ci/build_artifacts.sh "${{ inputs.artifact }}"
       - name: Upload artifacts to a GCS bucket (non-Windows runs)
         if: >-

.github/workflows/ci-build.yaml

@@ -75,7 +75,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
-        pip install uv
+        pip install uv~=0.5.30
         uv pip install --system .[minimum-jaxlib] -r build/test-requirements.txt
 
     - name: Run tests
@@ -88,7 +88,7 @@ jobs:
         JAX_SKIP_SLOW_TESTS: true
         PY_COLORS: 1
       run: |
-        pip install -e .
+        uv pip install --system -e .
         echo "JAX_NUM_GENERATED_CASES=$JAX_NUM_GENERATED_CASES"
         echo "JAX_ENABLE_X64=$JAX_ENABLE_X64"
         echo "JAX_ENABLE_CUSTOM_PRNG=$JAX_ENABLE_CUSTOM_PRNG"
@@ -113,7 +113,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
-        pip install uv
+        pip install uv~=0.5.30
         uv pip install --system -r docs/requirements.txt
     - name: Test documentation
       env:
@@ -147,7 +147,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
-        pip install uv
+        pip install uv~=0.5.30
         uv pip install --system -r docs/requirements.txt
     - name: Render documentation
       run: |
@@ -173,8 +173,9 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
-        pip install uv
-        uv pip install --system .[minimum-jaxlib] tensorflow -r build/test-requirements.txt
+        pip install uv~=0.5.30
+        uv pip install --system .[minimum-jaxlib] -r build/test-requirements.txt
+        uv pip install --system --pre tensorflow==2.19.0rc0
 
     - name: Run tests
       env:
@@ -184,7 +185,7 @@ jobs:
         JAX_SKIP_SLOW_TESTS: true
         PY_COLORS: 1
       run: |
-        pip install -e .
+        uv pip install --system -e .
         echo "JAX_NUM_GENERATED_CASES=$JAX_NUM_GENERATED_CASES"
         echo "JAX_ENABLE_X64=$JAX_ENABLE_X64"
         echo "JAX_ENABLE_CHECKS=$JAX_ENABLE_CHECKS"
@@ -205,7 +206,7 @@ jobs:
         python-version: 3.12
     - name: Install JAX
       run: |
-        pip install uv
+        pip install uv~=0.5.30
         uv pip install --system .[cuda12]
     - name: Build and install example project
       run: uv pip install --system ./examples/ffi[test]

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

@@ -59,42 +59,38 @@ jobs:
           git config --global --add safe.directory "$GITHUB_WORKSPACE"
       - 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
+          $PYTHON -m uv pip install -U -r build/test-requirements.txt -r build/collect-profile-requirements.txt
       - name: Install JAX
         run: |
-          $PYTHON -m pip uninstall -y jax jaxlib libtpu
+          $PYTHON -m uv pip uninstall jax jaxlib libtpu
           if [ "${{ matrix.jaxlib-version }}" == "head" ]; then
             # Build and install jaxlib at head
             $PYTHON build/build.py build --wheels=jaxlib \
                     --bazel_options=--config=rbe_linux_x86_64 \
                     --local_xla_path="$(pwd)/xla" \
                     --verbose
 
-            $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
+            # Install jaxlib, "jax" at head, and libtpu
+            $PYTHON -m uv pip install dist/*.whl \
+               -U -e . \
+               --pre libtpu -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
           elif [ "${{ matrix.jaxlib-version }}" == "pypi_latest" ]; then
-            $PYTHON -m pip install .[tpu] \
+            $PYTHON -m uv pip install .[tpu] \
               -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
 
           elif [ "${{ matrix.jaxlib-version }}" == "nightly" ]; then
-            $PYTHON -m pip install --pre . -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
-            $PYTHON -m pip install --pre libtpu \
-              -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-            $PYTHON -m pip install requests
+            $PYTHON -m uv pip install \
+              --pre . -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
+              libtpu -f https://storage.googleapis.com/jax-releases/libtpu_releases.html \
+              requests
 
           elif [ "${{ matrix.jaxlib-version }}" == "nightly+oldest_supported_libtpu" ]; then
-            $PYTHON -m pip install --pre . -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
             # TODO(phawkins): switch to libtpu, when the oldest release we support is a libtpu release.
-            $PYTHON -m pip install --pre libtpu-nightly==0.1.dev${{ env.LIBTPU_OLDEST_VERSION_DATE }} \
-              -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-            $PYTHON -m pip install requests
+            $PYTHON -m uv pip install \
+              --pre . -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
+              libtpu-nightly==0.1.dev${{ env.LIBTPU_OLDEST_VERSION_DATE }} \
+              -f https://storage.googleapis.com/jax-releases/libtpu_releases.html \
+              requests
           else
             echo "Unknown jaxlib-version: ${{ matrix.jaxlib-version }}"
             exit 1

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

@@ -74,8 +74,7 @@ jobs:
           git config --global --add safe.directory "$GITHUB_WORKSPACE"
       - name: Install JAX test requirements
         run: |
-          $JAXCI_PYTHON -m pip install -U -r build/test-requirements.txt
-          $JAXCI_PYTHON -m pip install -U -r build/collect-profile-requirements.txt
+          $JAXCI_PYTHON -m uv pip install -U -r build/test-requirements.txt -r build/collect-profile-requirements.txt
       - name: Build jaxlib at head with latest XLA
         run: |
           # Build and install jaxlib at head
@@ -86,7 +85,7 @@ jobs:
                   --verbose
 
           # Install libtpu
-          $JAXCI_PYTHON -m pip install --pre libtpu \
+          $JAXCI_PYTHON -m uv pip install --pre libtpu \
             -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
       # Halt for testing
       - name: Wait For Connection

.github/workflows/jax-array-api.yml

@@ -37,9 +37,8 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
-        pip install uv
-        uv pip install --system .[ci]
-        uv pip install --system pytest-xdist -r array-api-tests/requirements.txt
+        pip install uv~=0.5.30
+        uv pip install --system .[ci] pytest-xdist -r array-api-tests/requirements.txt
     - name: Run the test suite
       env:
         ARRAY_API_TESTS_MODULE: jax.numpy

.github/workflows/metal_plugin_ci.yml

@@ -35,15 +35,14 @@ jobs:
           rm -rf ${GITHUB_WORKSPACE}/jax-metal-venv
           python3 -m venv ${GITHUB_WORKSPACE}/jax-metal-venv
           source  ${GITHUB_WORKSPACE}/jax-metal-venv/bin/activate
-          pip install -U pip numpy wheel
-          pip install absl-py pytest
+          pip install uv~=0.5.30
+          uv pip install -U pip numpy wheel absl-py pytest
           if [[ "${{ matrix.jaxlib-version }}" == "nightly" ]]; then
-            pip install --pre jaxlib \
+            uv pip install --pre jaxlib \
               -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
           fi;
           cd jax
-          pip install .
-          pip install jax-metal
+          uv pip install . jax-metal
       - name: Run test
         run: |
           source  ${GITHUB_WORKSPACE}/jax-metal-venv/bin/activate