Add 3.13 and 3.13t to tpu matrix strategy and run only a single build for oldest supported libtpu

The following builds are run for TPU in the nightly/release workflow:
* v4-8: Python 3.10
* v5e-8: Python 3.13, Python 3.13t
* v6e-8: Python 3.11, Python 3.12

The following builds are run for TPU in the continuous workflow:
* v4-8: Python 3.10, libtpu=nightly
* v5e-8: Python 3.10, libtpu=nightly, oldest_supported_libtpu
* v6e-8: Python 3.10, libtpu=nightly

PiperOrigin-RevId: 752686662

Author: nitins17

.github/workflows/wheel_tests_continuous.yml

@@ -174,14 +174,24 @@ jobs:
           tpu-specs: [
             # {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"}
+            {type: "v5e-8", cores: "8", runner: "linux-x86-ct5lp-224-8tpu"},
+            {type: "v6e-8", cores: "8", runner: "linux-x86-ct6e-180-8tpu"}
           ]
+          libtpu-version-type: ["nightly", "oldest_supported_libtpu"]
+          exclude:
+            # Run a single config for oldest_supported_libtpu
+            - libtpu-version-type: "oldest_supported_libtpu"
+              tpu-specs:
+                type: "v4-8"
+            - libtpu-version-type: "oldest_supported_libtpu"
+              tpu-specs:
+                type: "v6e-8"
     name: "Pytest TPU (JAX artifacts version =  ${{ format('{0}', 'head') }})"
     with:
       runner: ${{ matrix.tpu-specs.runner }}
       cores: ${{ matrix.tpu-specs.cores }}
       tpu-type: ${{ matrix.tpu-specs.type }}
       python: ${{ matrix.python }}
       run-full-tpu-test-suite: "1"
-      libtpu-version-type: "nightly"
+      libtpu-version-type: ${{ matrix.libtpu-version-type }}
       gcs_download_uri: ${{ needs.build-jaxlib-artifact.outputs.gcs_upload_uri }}

.github/workflows/wheel_tests_nightly_release.yml

@@ -83,37 +83,50 @@ jobs:
     strategy:
         fail-fast: false # don't cancel all jobs on failure
         matrix:
-          # Skip Python 3.13 as it fails due to missing TensorFlow wheels (used for
-          # profiler_test.py, build/collect-profile-requirements.txt) for that version (b/402590302)
-          python: ["3.10", "3.11", "3.12"]
+          python: ["3.10", "3.11", "3.12", "3.13", "3.13-nogil"]
           tpu-specs: [
             # {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"},
             {type: "v6e-8", cores: "8", runner: "linux-x86-ct6e-180-8tpu"}
           ]
-          libtpu-version-type: ["pypi_latest", "nightly", "oldest_supported_libtpu"]
+          libtpu-version-type: ["pypi_latest", "nightly"]
           exclude:
+            # Exclude nightly for releases
             - libtpu-version-type: ${{ startsWith(github.ref_name, 'release/') && 'nightly' }}
-            - libtpu-version-type: ${{ startsWith(github.ref_name, 'release/') && 'oldest_supported_libtpu' }}
+            # Exclude pypi_latest for nightly releases
             - libtpu-version-type: ${{ !startsWith(github.ref_name, 'release/') && 'pypi_latest' }}
-            # Run a single Python version for v4-8 and v6e-8
+            # Run a single Python version for v4-8
             - tpu-specs:
                 type: "v4-8"
               python: "3.10"
             - tpu-specs:
                 type: "v4-8"
               python: "3.11"
+            - tpu-specs:
+                type: "v4-8"
+              python: "3.12"
+            - tpu-specs:
+                type: "v4-8"
+              python: "3.13-nogil"
+            # Run Python versions in between min and max for v6e-8
             - tpu-specs:
                 type: "v6e-8"
               python: "3.10"
             - tpu-specs:
                 type: "v6e-8"
-              python: "3.11"
+              python: "3.13"
+            - tpu-specs:
+                type: "v6e-8"
+              python: "3.13-nogil"
             # Run min and max Python versions for v5e-8
             - tpu-specs:
                 type: "v5e-8"
               python: "3.11"
+            - tpu-specs:
+                type: "v5e-8"
+              python: "3.12"
+
     name: "Pytest TPU (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})"
     with:
       runner: ${{ matrix.tpu-specs.runner }}

build/collect-profile-requirements.txt

@@ -1,4 +1,5 @@
-tensorflow
+# TF hasn't released 3.13 wheels yet (b/402590302)
+tensorflow; python_version<"3.13"
 tensorboard-plugin-profile<=2.19.0
 # Needed for the profile plugin to work without error
 protobuf

ci/run_pytest_tpu.sh

@@ -41,7 +41,9 @@ echo "Installed packages:"
 "$JAXCI_PYTHON" -c 'import sys; print("python version:", sys.version)'
 "$JAXCI_PYTHON" -c 'import jax; print("jax version:", jax.__version__)'
 "$JAXCI_PYTHON" -c 'import jaxlib; print("jaxlib version:", jaxlib.__version__)'
-strings /usr/local/lib/"$JAXCI_PYTHON"/dist-packages/libtpu/libtpu.so | grep 'Built on'
+# Free-threaded builds use "-nogil" as the suffix for the binary and "t" for its
+# dist-packages path
+strings /usr/local/lib/"${JAXCI_PYTHON//-nogil/t}"/dist-packages/libtpu/libtpu.so | grep 'Built on'
 "$JAXCI_PYTHON" -c 'import jax; print("libtpu version:",jax.lib.xla_bridge.get_backend().platform_version)'
 
 # Set up all common test environment variables