diff --git a/.github/workflows/wheel_tests_continuous.yml b/.github/workflows/wheel_tests_continuous.yml index 07e2eb4fd978..aee1f3311d98 100644 --- a/.github/workflows/wheel_tests_continuous.yml +++ b/.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') }})" diff --git a/.github/workflows/wheel_tests_nightly_release.yml b/.github/workflows/wheel_tests_nightly_release.yml index 8a5703412936..7a769b7adee4 100644 --- a/.github/workflows/wheel_tests_nightly_release.yml +++ b/.github/workflows/wheel_tests_nightly_release.yml @@ -200,6 +200,13 @@ jobs: - tpu-specs: type: "v5e-8" python: "3.12" + # Run min and max Python versions for v7x-8 + - tpu-specs: + type: "v7x-8" + python: "3.11" + - 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 +230,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: @@ -247,6 +255,13 @@ jobs: - tpu-specs: type: "v5e-8" python: "3.12" + # Run min and max Python versions for v7x-8 + - tpu-specs: + type: "v7x-8" + python: "3.11" + - tpu-specs: + type: "v7x-8" + python: "3.12" name: "Bazel tests TPU (JAX artifacts version = ${{ startsWith(github.ref_name, 'release/') && 'latest release' || 'nightly' }})" with: