Add linux python 3.13t nightly tests

* Python wheels follow a naming convention: standard wheels use the pattern `*-cp<python_version>-cp<python_version>-*`, while free-threaded wheels use `*-cp<python_version>-cp<python_version>t-*`. Update the pytest workflows to look for free-threaded wheels and ensure that standard wheel tests exclude free-threaded wheels.

* Skip zstandard for python3.13-nogil due to compilation failure indygreg/python-zstandard#231.

PiperOrigin-RevId: 732070585

.github/workflows/pytest_cpu.yml

@@ -79,11 +79,15 @@ jobs:
 
           # Get the major and minor version of Python.
           # E.g if JAXCI_HERMETIC_PYTHON_VERSION=3.10, then python_major_minor=310
-          python_major_minor=$(echo "$JAXCI_HERMETIC_PYTHON_VERSION" | tr -d '.')
+          # E.g if JAXCI_HERMETIC_PYTHON_VERSION=3.13-nogil, then python_major_minor=313t
+          python_major_minor=$(echo "${JAXCI_HERMETIC_PYTHON_VERSION//-nogil/t}" | tr -d '.')
 
           echo "OS=${os}" >> $GITHUB_ENV
           echo "ARCH=${arch}" >> $GITHUB_ENV
-          echo "PYTHON_MAJOR_MINOR=${python_major_minor}" >> $GITHUB_ENV
+          # Python wheels follow a naming convention: standard wheels use the pattern
+          # `*-cp<py_version>-cp<py_version>-*`, while free-threaded wheels use
+          # `*-cp<py_version>-cp<py_version>t-*`.
+          echo "PYTHON_MAJOR_MINOR=cp${python_major_minor%t}-cp${python_major_minor}-" >> $GITHUB_ENV
       - name: Download jaxlib wheel from GCS (non-Windows runs)
         id: download-wheel-artifacts-nw
         # Set continue-on-error to true to prevent actions from failing the workflow if this step
@@ -128,7 +132,16 @@ jobs:
         run: |
           # TODO(srnitin): Remove after uv is installed in the Windows Dockerfile
           $JAXCI_PYTHON -m pip install uv~=0.5.30
-          $JAXCI_PYTHON -m uv pip install -r build/requirements.in
+          # python 3.13t cannot compile zstandard 0.23.0 due to
+          # /~https://github.com/indygreg/python-zstandard/issues/231. Remove this once zstandard
+          # has a prebuilt wheel for 3.13t or an env marker is available for free threading python
+          # in requirements.in.
+          if [[ $JAXCI_PYTHON =~ "python3.13-nogil" ]]; then
+            grep -v "zstandard" build/requirements.in > build/requirements_without_zstandard.txt
+            $JAXCI_PYTHON -m uv pip install -r build/requirements_without_zstandard.txt
+          else
+            $JAXCI_PYTHON -m uv pip install -r build/requirements.in
+          fi
       # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main

.github/workflows/pytest_cuda.yml

@@ -77,11 +77,15 @@ jobs:
 
           # Get the major and minor version of Python.
           # E.g if JAXCI_HERMETIC_PYTHON_VERSION=3.10, then python_major_minor=310
-          python_major_minor=$(echo "$JAXCI_HERMETIC_PYTHON_VERSION" | tr -d '.')
+          # E.g if JAXCI_HERMETIC_PYTHON_VERSION=3.13-nogil, then python_major_minor=313t
+          python_major_minor=$(echo "${JAXCI_HERMETIC_PYTHON_VERSION//-nogil/t}" | tr -d '.')
 
           echo "OS=${os}" >> $GITHUB_ENV
           echo "ARCH=${arch}" >> $GITHUB_ENV
-          echo "PYTHON_MAJOR_MINOR=${python_major_minor}" >> $GITHUB_ENV
+          # Python wheels follow a naming convention: standard wheels use the pattern
+          # `*-cp<py_version>-cp<py_version>-*`, while free-threaded wheels use
+          # `*-cp<py_version>-cp<py_version>t-*`.
+          echo "PYTHON_MAJOR_MINOR=cp${python_major_minor%t}-cp${python_major_minor}-" >> $GITHUB_ENV
       - name: Download the wheel artifacts from GCS
         id: download-wheel-artifacts
         # Set continue-on-error to true to prevent actions from failing the workflow if this step
@@ -106,7 +110,17 @@ jobs:
           echo "Skipping the test run."
           exit 1
       - name: Install Python dependencies
-        run: $JAXCI_PYTHON -m uv pip install -r build/requirements.in
+        run: |
+          # python 3.13t cannot compile zstandard 0.23.0 due to
+          # /~https://github.com/indygreg/python-zstandard/issues/231. Remove this once zstandard
+          # has a prebuilt wheel for 3.13t or an env marker is available for free threading python
+          # in requirements.in.
+          if [[ $JAXCI_PYTHON =~ "python3.13-nogil" ]]; then
+            grep -v "zstandard" build/requirements.in > build/requirements_without_zstandard.txt
+            $JAXCI_PYTHON -m uv pip install -r build/requirements_without_zstandard.txt
+          else
+            $JAXCI_PYTHON -m uv pip install -r build/requirements.in
+          fi
       # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main

.github/workflows/wheel_tests_nightly_release.yml

@@ -31,8 +31,11 @@ jobs:
           # Runner OS and Python values need to match the matrix stategy of our internal CI jobs
           # that build the wheels.
           runner: ["linux-x86-n2-64", "linux-arm64-t2a-48", "windows-x86-n2-64"]
-          python: ["3.10","3.11", "3.12", "3.13"]
+          python: ["3.10","3.11", "3.12", "3.13", "3.13-nogil"]
           enable-x64: [0]
+          exclude:
+            - runner: "windows-x86-n2-64"
+              python: "3.13-nogil"
     with:
       runner: ${{ matrix.runner }}
       python: ${{ matrix.python }}
@@ -50,7 +53,7 @@ jobs:
           # Runner OS and Python values need to match the matrix stategy of our internal CI jobs
           # that build the wheels.
           runner: ["linux-x86-g2-48-l4-4gpu"]
-          python: ["3.10","3.11", "3.12", "3.13"]
+          python: ["3.10","3.11", "3.12", "3.13", "3.13-nogil"]
           cuda: ["12.3", "12.1"]
           enable-x64: [0]
     with:

jaxlib/jax.bzl

@@ -64,7 +64,7 @@ PLATFORM_TAGS_DICT = {
 
 # TODO(vam): remove this once zstandard builds against Python 3.13
 def get_zstandard():
-    if HERMETIC_PYTHON_VERSION == "3.13":
+    if HERMETIC_PYTHON_VERSION == "3.13" or HERMETIC_PYTHON_VERSION == "3.13-ft":
         return []
     return ["@pypi_zstandard//:pkg"]