Revert "T5X arm64: Pin grain 0.0.4 source (#520)" (#537)

This reverts commit 6f7f5dd.

Now airio pin grain-nightly to 0.0.6:
https://github.com/google/airio/blob/main/setup.py#L52

.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

.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