Merge pull request #2380 from AI-Hypercomputer:tunix_name

PiperOrigin-RevId: 814731290

Author: Google-ML-Automation

maxtext_jax_ai_image.Dockerfile

@@ -47,9 +47,9 @@ RUN if [ "$DEVICE" = "tpu" ] && [ "$JAX_STABLE_STACK_BASEIMAGE" = "us-docker.pkg
         python3 -m pip install -r /deps/requirements_with_jax_ai_image.txt; \
   fi
 
-# Install tunix at a pinned commit for TPU devices, skip for GPU
+# Install google-tunix for TPU devices, skip for GPU
 RUN if [ "$DEVICE" = "tpu" ]; then \
-        python3 -m pip install 'google-tunix @ https://github.com/google/tunix/archive/6c2a613217ed6bb1c9d60088d6e3c67184821c69.zip'; \
+        python3 -m pip install 'google-tunix>=0.1.0'; \
   fi
 
 # Now copy the remaining code (source files that may change frequently)

pyproject.toml

@@ -90,6 +90,7 @@ dependencies = [
     "google-genai>=1.33.0",
     "google-pasta>=0.2.0",
     "google-resumable-media>=2.7.2",
+    "google-tunix>=0.1.0",
     "googleapis-common-protos>=1.70.0",
     "grain>=0.2.12",
     "grpc-google-iam-v1>=0.14.2",
@@ -193,6 +194,7 @@ dependencies = [
     "pytype>=2024.10.11",
     "pytz>=2025.2",
     "pyyaml>=6.0.2",
+    "qwix>=0.1.0",
     "regex>=2025.9.1",
     "requests-oauthlib>=2.0.0",
     "requests>=2.32.5",

requirements.txt

@@ -29,6 +29,7 @@ pyink
 pylint
 pytest
 pytype
+qwix
 sentencepiece
 tensorboard-plugin-profile
 tensorboardx
@@ -39,4 +40,3 @@ tiktoken
 transformers
 google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/daedc21c393f23449fb54ddc4f75fca34348ea9c.zip
 mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip
-qwix @ https://github.com/google/qwix/archive/f2fd7b9114ff8d09e5b0131a453351578502da8a.zip

requirements_with_jax_ai_image.txt

@@ -18,9 +18,9 @@ pyink
 pylint
 pytest
 pytype
+qwix
 sentencepiece>=0.2.0
 tensorflow-datasets
 tensorflow-text>=2.17.0
 tiktoken
 transformers
-qwix @ https://github.com/google/qwix/archive/f2fd7b9114ff8d09e5b0131a453351578502da8a.zip

setup.sh

@@ -190,8 +190,8 @@ if [[ "$MODE" == "stable" || ! -v MODE ]]; then
         fi
 
         # TODO: Once tunix has support for GPUs, move it from here to requirements.txt
-        echo "Installing tunix for stable TPU environment"
-        python3 -m uv pip install 'google-tunix @ https://github.com/google/tunix/archive/6c2a613217ed6bb1c9d60088d6e3c67184821c69.zip'
+        echo "Installing google-tunix for stable TPU environment"
+        python3 -m uv pip install 'google-tunix>=0.1.0'
 
         if [[ -n "$LIBTPU_GCS_PATH" ]]; then
             # Install custom libtpu

src/install_maxtext_extra_deps/extra_deps_from_github.txt

@@ -1,6 +1,3 @@
 datasets @ https://github.com/huggingface/datasets/archive/6790e138c00b87a1ddc72184f89e7814cf784360.zip
 google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/daedc21c393f23449fb54ddc4f75fca34348ea9c.zip
 mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip
-qwix @ https://github.com/google/qwix/archive/f2fd7b9114ff8d09e5b0131a453351578502da8a.zip
-google-tunix @ https://github.com/google/tunix/archive/6c2a613217ed6bb1c9d60088d6e3c67184821c69.zip
-