Fix for linux-x86-ct5lp-224-8tpu

Signed-off-by: Vladimir Suvorov <suvorovv@google.com>

Author: mydatascience

.github/workflows/tpu-tests.yml

@@ -35,7 +35,10 @@ jobs:
     runs-on: [linux-x86-ct5lp-224-8tpu]
     container:
       image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/tpu:latest
-      options: --privileged
+      options: --privileged --cpus=4 --memory=8Gi
+      env:
+        TPU_ACCELERATOR_TYPE: ""
+        JAX_PLATFORMS: "tpu"
     steps:
     - name: Test gsutil installation
       run: which gsutil >/dev/null 2>&1 || { echo >&2 "gsutil is required but not installed. Aborting"; exit 24;}
@@ -46,7 +49,7 @@ jobs:
     needs: prelim
     runs-on: [linux-x86-ct5lp-224-8tpu]
     container:
-      image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/tpu:latest
+      image: python:3.12-slim
       options: --privileged
     steps:
     - name: Checkout code
@@ -126,7 +129,7 @@ jobs:
     needs: prelim
     runs-on: [linux-x86-ct5lp-224-8tpu]
     container:
-      image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/tpu:latest
+      image: python:3.12-slim
       options: --privileged
     steps:
     - name: Checkout code