Internal CI change.

PiperOrigin-RevId: 718511240

Author: cantonios

.github/workflows/build_and_test.yml

@@ -9,19 +9,15 @@ on:
   pull_request:
 
 jobs:
-  build_and_test:
+  build-and-test-cpu:
     runs-on: ubuntu-24.04
-    strategy:
-      matrix:
-        python-version: ["3.10"]
-
     steps:
       - uses: actions/checkout@v4
 
-      - name: Set up Python ${{ matrix.python-version }}
+      - name: Set up Python 3.10
         uses: actions/setup-python@v5
         with:
-          python-version: ${{ matrix.python-version }}
+          python-version: '3.10'
           cache: 'pip'
 
       - name: Display Python version
@@ -42,26 +38,26 @@ jobs:
         uses: actions/cache@v4
         with:
           path: "/home/runner/.cache/bazel"
-          key: bazel-${{ matrix.python-version }}-${{ github.ref_name }}-${{ github.sha }}
+          key: bazel-py3.10-${{ github.ref_name }}-${{ github.sha }}
           restore-keys: |
-            bazel-${{ matrix.python-version }}-${{ github.ref_name }}
-            bazel-${{ matrix.python-version }}-
+            bazel-py3.10-${{ github.ref_name }}
+            bazel-py3.10-
             bazel-
 
       - if: github.event_name == 'pull_request'
         name: Mount bazel cache (pull-request)
         uses: actions/cache/restore@v4
         with:
           path: "/home/runner/.cache/bazel"
-          key: bazel-${{ matrix.python-version }}-${{ github.base_ref }}-${{ github.event.pull_request.base.sha }}
+          key: bazel-py3.10-${{ github.base_ref }}-${{ github.event.pull_request.base.sha }}
           restore-keys: |
-            bazel-${{ matrix.python-version }}-${{ github.base_ref }}
-            bazel-${{ matrix.python-version }}-
+            bazel-py3.10-${{ github.base_ref }}
+            bazel-py3.10-
             bazel-
 
       - name: Build all targets
         run: |
-          export HERMETIC_PYTHON_VERSION=${{ matrix.python-version }}
+          export HERMETIC_PYTHON_VERSION=3.10
           bazel build //...
 
       - name: Build pip wheel

README.md

@@ -1,6 +1,7 @@
 # jax_tpu_embedding
 
-[![Unittests](https://github.com/jax-ml/jax_tpu_embedding/actions/workflows/pytest_and_autopublish.yml/badge.svg)](https://github.com/jax-ml/jax_tpu_embedding/actions/workflows/pytest_and_autopublish.yml)
-[![PyPI version](https://badge.fury.io/py/jax_tpu_embedding.svg)](https://badge.fury.io/py/jax_tpu_embedding)
+[![Unittests](https://github.com/jax-ml/jax-tpu-embedding/actions/workflows/build_and_test.yml/badge.svg)](https://github.com/jax-ml/jax-tpu-embedding/actions/workflows/build_and_test.yml)
+
+Usage instructions coming soon!
 
 *This is not an officially supported Google product.*