Add a jaxlib at head build to the cloud-tpu-ci-nightly workflow

This will allow us to test TPU compatibility with jaxlib at head. Also, enable v4 runners as they are now online.

PiperOrigin-RevId: 699155667

Author: nitins17

.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