fix release workflow

PiperOrigin-RevId: 801894880

Author: hbq1

.github/workflows/ci.yml

@@ -2,9 +2,9 @@ name: ci
 
 on:
   push:
-    branches: ["master"]
+    branches: ["main"]
   pull_request:
-    branches: ["master"]
+    branches: ["main"]
   schedule:
     - cron: '30 2 * * *'
 
@@ -18,7 +18,7 @@ concurrency:
 
 jobs:
   build-and-test:
-    name: "Python ${{ matrix.python-version }} on ${{ matrix.os }} jax=${{ matrix.jax-version}}"
+    name: "Pytest ${{ matrix.python-version }} on ${{ matrix.os }} jax=${{ matrix.jax-version}}"
     runs-on: "${{ matrix.os }}"
 
     strategy:
@@ -32,12 +32,10 @@ jobs:
             jax-version: "nightly"
 
     steps:
-      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
-      - uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b  # v5.3.0
+      - uses: actions/checkout@v5
+      - uses: actions/setup-python@v5.3.0
         with:
             python-version: "${{ matrix.python-version }}"
-            cache: "pip"
-            cache-dependency-path: '**/requirements*.txt'
       - name: Run CI tests
         run: JAX_VERSION="${{ matrix.jax-version }}" bash test.sh
         shell: bash

.github/workflows/pypi-publish.yml

@@ -18,9 +18,10 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install setuptools wheel twine
+        pip install setuptools wheel twine build
     - name: Check consistency between the package version and release tag
       run: |
+        pip install .
         RELEASE_VER=${GITHUB_REF#refs/*/}
         PACKAGE_VER="v`python -c 'import chex; print(chex.__version__)'`"
         if [ $RELEASE_VER != $PACKAGE_VER ]
@@ -29,6 +30,6 @@ jobs:
         fi
     - name: Build
       run: |
-        python setup.py sdist bdist_wheel
+        python -m build
     - name: Publish package distributions to PyPI
       uses: pypa/gh-action-pypi-publish@release/v1