[CI] Modify TPU v7x jobs to include bazel, and exclude some python versions

PiperOrigin-RevId: 845803404

Author: MichaelHudgins

.github/workflows/wheel_tests_continuous.yml

@@ -268,6 +268,7 @@ jobs:
           tpu-specs: [
             {type: "v4-8", cores: "4", runner: "linux-x86-ct4p-240-4tpu"},
             {type: "v5e-8", cores: "8", runner: "linux-x86-ct5lp-224-8tpu"},
+            {type: "v7x-8", cores: "8", runner: "linux-x86-tpu7x-224-4tpu"},
           ]
           libtpu-version-type: ["nightly"]
     name: "Bazel tests TPU (JAX artifacts version =  ${{ format('{0}', 'head') }})"

.github/workflows/wheel_tests_nightly_release.yml

@@ -193,13 +193,17 @@ jobs:
             - tpu-specs:
                 type: "v6e-8"
               python: "3.13-nogil"
-            # Run min and max Python versions for v5e-8
+            # Run max Python versions for v5e-8
             - tpu-specs:
                 type: "v5e-8"
               python: "3.11"
             - tpu-specs:
                 type: "v5e-8"
               python: "3.12"
+            # Run min and max Python versions for v7x-8
+            - tpu-specs:
+                type: "v7x-8"
+              python: "3.12"
 
     name: "Pytest TPU (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})"
     with:
@@ -223,6 +227,7 @@ jobs:
             # {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: "v7x-8", cores: "8", runner: "linux-x86-tpu7x-224-4tpu"},
           ]
           libtpu-version-type: ["pypi_latest", "nightly"]
           exclude:
@@ -240,12 +245,13 @@ jobs:
             - tpu-specs:
                 type: "v4-8"
               python: "3.13-nogil"
-            # Run min and max Python versions for v5e-8
+            # Run max Python versions for v5e-8
             - tpu-specs:
                 type: "v5e-8"
-              python: "3.11"
+              python: "3.12"
+            # Run min and max Python versions for v7x-8
             - tpu-specs:
-                type: "v5e-8"
+                type: "v7x-8"
               python: "3.12"
 
     name: "Bazel tests TPU (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})"