Merge branch 'main' into srnitin/test-scripts-workflows

Author: nitins17

.github/workflows/pytest_cpu.yml

@@ -73,11 +73,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)
         if: ${{ !contains(inputs.runner, 'windows-x86') }}
         run: >-

.github/workflows/pytest_tpu.yml

@@ -36,8 +36,13 @@ 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: ["linux-arm64-t2a-48"]
+              python: ["3.13-nogil"]
+            - runner: ["windows-x86-n2-64"]
+              python: ["3.13-nogil"]
     with:
       runner: ${{ matrix.runner }}
       python: ${{ matrix.python }}

.github/workflows/wait-for-connection-test.yaml

@@ -16,7 +16,10 @@ scipy>=1.13.1
 
 ml_dtypes>=0.4.0
 opt_einsum
-zstandard
+# python 3.13t cannot compile zstandard 0.23.0 due to
+# https://github.com/indygreg/python-zstandard/issues/231.
+# pip cannot differentiate between python 3.13 and 3.13t.
+zstandard; python_version<="3.12"
 etils[epath]
 # TODO(ybaturina): remove setuptools version
 setuptools<71.0.0