update bazel cpu workflow

Author: nitins17

.github/workflows/bazel_cpu.yml

@@ -13,8 +13,8 @@ jobs:
         runner: ["linux-x86-n2-64", "linux-arm64-t2a-48"]
 
     runs-on: ${{ matrix.runner }}
-    container: ${{ (contains(matrix.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/build:670606426-python3.9') ||
-                   (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/build-arm64:jax-') }}
+    container: ${{ (contains(matrix.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/build:670606426-python3.12') ||
+                   (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/linux-arm64-arc-container:latest') }}
 
     env:
       # Disable running `run_docker_container.sh`. JAX self-hosted runners already run in a Docker
@@ -28,6 +28,4 @@ jobs:
       - name: Wait For Connection
         uses: ./actions/ci_connection/
       - name: Run Bazel CPU Tests
-        env:
-          JAXCI_ENV_FILE: ci/envs/run_tests/bazel_cpu
-        run: ./ci/run_bazel_test.sh
+        run: ./ci/run_bazel_test.sh "ci/envs/run_tests/bazel_cpu"