Add tests on python 3.13

PiperOrigin-RevId: 749051219

Author: vroulet

.github/workflows/tests.yml

@@ -26,7 +26,7 @@ jobs:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
       - uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
         with:
-            python-version: "3.11"
+            python-version: "3.12"
             cache: "pip"
             cache-dependency-path: "pyproject.toml"
       - name: Install linting dependencies
@@ -51,7 +51,7 @@ jobs:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
       - uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
         with:
-            python-version: "3.11"
+            python-version: "3.12"
             cache: "pip"
             cache-dependency-path: "pyproject.toml"
       - name: Install ruff and lint check
@@ -65,13 +65,27 @@ jobs:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
       - uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55 # v5.5.0
       - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1
+  pytype:
+    name: "Check types with pytype"
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+      - uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
+        with:
+            python-version: "3.13"
+            cache: "pip"
+            cache-dependency-path: "pyproject.toml"
+      - name: Check types with pytype
+        run: |
+          python3 -m pip install --quiet --upgrade pytype
+          pytype "optax" -j auto --keep-going --disable import-error
   doctests:
     needs: [pre-commit, linting, ruff-lint]  # do not run doctests if linting fails
     name: "Doctests on ${{ matrix.os }} with Python ${{ matrix.python-version }}"
     runs-on: "${{ matrix.os }}"
     strategy:
       matrix:
-        python-version: ["3.11"]  # only build docs with a somewhat latest python
+        python-version: ["3.12"]  # only build docs with a somewhat latest python
         os: [ubuntu-latest]
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
@@ -93,14 +107,14 @@ jobs:
     runs-on: "${{ matrix.os }}"
     strategy:
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
         os: [ubuntu-latest]
         jax-version: [newest]
         include:
           - python-version: "3.10"
             os: "ubuntu-latest"
             jax-version: "0.4.27"  # Keep version in sync with pyproject.toml and copy.bara.sky!
-          - python-version: "3.12"
+          - python-version: "3.13"
             os: "ubuntu-latest"
             jax-version: "nightly"
     steps:

test.sh

@@ -69,7 +69,7 @@ python3 -m pip wheel --no-deps dist/optax-*.tar.gz --wheel-dir "${TEMP_DIR}"
 python3 -m pip install --quiet "${TEMP_DIR}/optax-"*.whl
 
 # Check types with pytype.
-python3 -m pip install --quiet pytype
+python3 -m pip install --quiet --upgrade pytype
 pytype "optax" -j auto --keep-going --disable import-error
 
 # Run tests using pytest.