From c318902ee5f1a66dfc94e445fdf93bec4ae6e2b1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fr=C3=A9d=C3=A9ric=20Bastien?= <fbastien@nvidia.com>
Date: Tue, 13 Feb 2024 20:25:56 -0500
Subject: [PATCH] Revert "T5X arm64: Pin grain 0.0.4 source (#520)" (#537)

This reverts commit 6f7f5dd5e048b285a2365ab68bee6774dbd80b29.

Now airio pin grain-nightly to 0.0.6:
https://github.com/google/airio/blob/main/setup.py#L52
---
 .github/container/Dockerfile.t5x.arm64 | 5 ++---
 .github/container/manifest.yaml        | 6 ++++++
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/.github/container/Dockerfile.t5x.arm64 b/.github/container/Dockerfile.t5x.arm64
index 4ee584883..f1f6c888d 100644
--- a/.github/container/Dockerfile.t5x.arm64
+++ b/.github/container/Dockerfile.t5x.arm64
@@ -120,13 +120,12 @@ COPY --from=array_record-builder /tmp/array_record/all_dist/array_record*linux_a
             tensorflow-datasets
 # END
 
+RUN get-source.sh -l grain -m ${MANIFEST_FILE}
+
 RUN <<"EOT" bash -exu
 set -o pipefail
 
-# Check out source of grain-nightly 0.0.4
-git clone https://github.com/google/grain ${SRC_PATH_GRAIN}
 pushd ${SRC_PATH_GRAIN}
-git checkout fa79b9dea81ffb00555a6c2ae2898be4bdd5e564
 
 # Make bazel stop complaining about sharding and disable some tests with missing bazel build deps
 sed -i 's| bazel test | bazel test --test_sharding_strategy=disabled |' ./grain/oss/build_whl.sh
diff --git a/.github/container/manifest.yaml b/.github/container/manifest.yaml
index 9eccb9768..50662e627 100644
--- a/.github/container/manifest.yaml
+++ b/.github/container/manifest.yaml
@@ -132,6 +132,12 @@ mujoco:
   tracking_ref: main
   latest_verified_commit: ced37ffbb237584512311b041bce3124e3b2cc2a
   mode: git-clone
+grain:
+  # Used only in ARM t5x builds
+  url: https://github.com/google/grain.git
+  tracking_ref: main
+  latest_verified_commit: 41f80fb540547ec3aefe7cec57fe43228560b1cf
+  mode: git-clone
 mujoco-mpc:
   url: https://github.com/google-deepmind/mujoco_mpc.git
   tracking_ref: main