MaxText: grain-nightly -> grain, t5x: make optax as in jax-ai-stack

NVIDIA · Dec 21, 2024 · 35455da · 35455da
1 parent cdd7330
commit 35455da
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/.github/container/Dockerfile.maxtext b/.github/container/Dockerfile.maxtext
@@ -67,6 +67,7 @@ for pattern in \
     "s|absl-py|absl-py>=2.1.0|g" \
     "s|protobuf==3.20.3|protobuf>=3.19.0|g" \
     "s|tensorflow-datasets|tensorflow-datasets>=4.8.0|g" \
+    "s|grain-nightly|grain|g" \
   ; do
     sed -i "${pattern}" ${SRC_PATH_MAXTEXT}/requirements.txt;
 done

diff --git a/.github/container/Dockerfile.t5x b/.github/container/Dockerfile.t5x
@@ -73,6 +73,7 @@ echo "seqio-nightly>=0.0.18.dev20240714" >> /opt/pip-tools.d/requirements-t5x.in
 # 2. Remove head-of-tree specs from select dependencies
 pushd ${SRC_PATH_T5X}
 sed -i "s| @ git+https://github.com/google/flax#egg=flax||g" setup.py
+sed -i "s| @ git+https://github.com/deepmind/optax#egg=optax||g" setup.py
 
 # for ARM64 build
 if [[ "$(dpkg --print-architecture)" == "arm64" ]]; then