From d04c2d1a7416db22c1531e3d9fe7857d52962815 Mon Sep 17 00:00:00 2001
From: hmonishN <143435143+hmonishN@users.noreply.github.com>
Date: Thu, 16 Nov 2023 16:43:55 -0800
Subject: [PATCH] Update _sandbox.yaml

---
 .github/workflows/_sandbox.yaml | 25 ++++++++++++-------------
 1 file changed, 12 insertions(+), 13 deletions(-)

diff --git a/.github/workflows/_sandbox.yaml b/.github/workflows/_sandbox.yaml
index 6b7adfecd..a57eec3ec 100644
--- a/.github/workflows/_sandbox.yaml
+++ b/.github/workflows/_sandbox.yaml
@@ -116,19 +116,19 @@ jobs:
       REF_XLA: ${{ needs.metadata.outputs.REF_XLA }}
     secrets: inherit
 
-  build-t5x:
+  build-te:
     needs: [metadata, build-jax]
-    uses: ./.github/workflows/_build_t5x.yaml
+    uses: ./.github/workflows/_build_te.yaml
     with:
       BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
       BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAGS }}
-      REPO_T5X: ${{ needs.metadata.outputs.REPO_T5X }}
-      REF_T5X: ${{ needs.metadata.outputs.REF_T5X }}
+      REPO_TE: ${{ needs.metadata.outputs.REPO_TE }}
+      REF_TE: ${{ needs.metadata.outputs.REF_TE }}
     secrets: inherit
 
 
   build-summary:
-    needs: [build-base, build-jax, build-t5x]
+    needs: [build-base, build-jax, build-te]
     # needs: [build-base, build-jax, build-te, build-t5x, build-pax, build-pax-aarch64, build-rosetta-t5x, build-rosetta-pax]
     if: always()
     runs-on: ubuntu-22.04
@@ -143,12 +143,12 @@ jobs:
           | ------------ | -------------------------------------------------- |
           | Base         | ${{ needs.build-base.outputs.DOCKER_TAGS }}        |
           | JAX          | ${{ needs.build-jax.outputs.DOCKER_TAGS }}         |
-          | T5X          | ${{ needs.build-t5x.outputs.DOCKER_TAGS }}         |
+          | JAX-TE       | ${{ needs.build-te.outputs.DOCKER_TAGS }}          |
 
           EOF
 
   retrofit-containers:
-    needs: [build-base, build-jax, build-t5x]
+    needs: [build-base, build-jax, build-te]
     if: always()
     runs-on: ubuntu-22.04
     env:
@@ -227,18 +227,17 @@ jobs:
       JAX_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAGS }}
     secrets: inherit
 
-  test-t5x:
-    needs: build-t5x
-    uses: ./.github/workflows/_test_t5x.yaml
+  test-te:
+    needs: build-te
+    uses: ./.github/workflows/_test_te.yaml
     with:
-      T5X_IMAGE: ${{ needs.build-t5x.outputs.DOCKER_TAGS }}
+      JAX_TE_IMAGE: ${{ needs.build-te.outputs.DOCKER_TAGS }}
     secrets: inherit
-
  
   finalize:
     if: always()
     # TODO: use dynamic matrix to make dependencies self-updating
-    needs: [build-summary, test-distribution, test-jax, test-t5x]
+    needs: [build-summary, test-distribution, test-jax, test-te]
     uses: ./.github/workflows/_finalize.yaml
     with:
       PUBLISH_BADGE: false