diff --git a/docs/guides/data_input_pipeline/data_input_tfds.md b/docs/guides/data_input_pipeline/data_input_tfds.md
index acbf064055..2b5bffedc0 100644
--- a/docs/guides/data_input_pipeline/data_input_tfds.md
+++ b/docs/guides/data_input_pipeline/data_input_tfds.md
@@ -1,4 +1,18 @@
-# TFDS pipeline
+# TFDS pipeline (Deprecated)
+
+```{warning}
+The TFDS input pipeline (`dataset_type=tfds`) is **deprecated**. We recommend migrating to the [Grain pipeline](data_input_grain.md) with `dataset_type=grain` and `grain_file_type=tfrecord`. You can keep the same TFRecord dataset paths.
+```
+
+````{note}
+TensorFlow and TensorFlow Datasets (TFDS) are optional dependencies in MaxText. If you need to use the legacy TFDS pipeline, install the optional dependencies by running:
+
+```bash
+install_tpu_pre_train_extra_deps --with-tf
+# or for GPU:
+# install_cuda12_pre_train_extra_deps --with-tf
+```
+````
 
 1. Download the Allenai C4 dataset in TFRecord format to a Cloud Storage bucket. For information about cost, see [this discussion](https://github.com/allenai/allennlp/discussions/5056)
 
diff --git a/docs/install_maxtext.md b/docs/install_maxtext.md
index d01bbc3fb6..6335b0c4e8 100644
--- a/docs/install_maxtext.md
+++ b/docs/install_maxtext.md
@@ -53,6 +53,8 @@ This is the easiest way to get started with the latest stable version.
      ```bash
      uv pip install maxtext[tpu]=={{version}} --resolution=lowest
      install_tpu_pre_train_extra_deps
+     # Pass --with-tf to install optional TensorFlow/TFDS and JetStream dependencies if needed:
+     # install_tpu_pre_train_extra_deps --with-tf
      ```
 
    - **Option 2:** Install `maxtext[cuda12]`, used for pre-training and decoding
@@ -61,6 +63,8 @@ This is the easiest way to get started with the latest stable version.
      ```bash
      uv pip install maxtext[cuda12]=={{version}} --resolution=lowest
      install_cuda12_pre_train_extra_deps
+     # Pass --with-tf to install optional TensorFlow/TFDS and JetStream dependencies if needed:
+     # install_cuda12_pre_train_extra_deps --with-tf
      ```
 
    - **Option 3:** Install `maxtext[tpu-post-train]`, used for post-training on
@@ -130,6 +134,8 @@ environment to avoid dependency conflicts.
      ```bash
      uv pip install -e .[tpu] --resolution=lowest
      install_tpu_pre_train_extra_deps
+     # Pass --with-tf to install optional TensorFlow/TFDS and JetStream dependencies if needed:
+     # install_tpu_pre_train_extra_deps --with-tf
      ```
 
    - **Option 2:** Install `.[cuda12]`
@@ -137,6 +143,8 @@ environment to avoid dependency conflicts.
      ```bash
      uv pip install -e .[cuda12] --resolution=lowest
      install_cuda12_pre_train_extra_deps
+     # Pass --with-tf to install optional TensorFlow/TFDS and JetStream dependencies if needed:
+     # install_cuda12_pre_train_extra_deps --with-tf
      ```
 
    - **Option 3:** Install `.[tpu-post-train]`
diff --git a/src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile b/src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
index 56c578cf4c..e290a03816 100644
--- a/src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
+++ b/src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
@@ -38,6 +38,10 @@ ENV ENV_JAX_VERSION=$JAX_VERSION
 ARG DEVICE
 ENV ENV_DEVICE=$DEVICE
 
+# TODO: remove default once separate TF and TF-free nightly image workflows are established
+ARG TF=true
+ENV ENV_TF=$TF
+
 ARG PACKAGE_DIR
 ENV PACKAGE_DIR=$PACKAGE_DIR
 
@@ -56,10 +60,10 @@ COPY ${PACKAGE_DIR}/dependencies/scripts/ src/dependencies/scripts/
 COPY ${PACKAGE_DIR}/maxtext/integration/vllm/ src/maxtext/integration/vllm/
 
 # Install dependencies - these steps are cached unless the copied files change
-RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION DEVICE=${ENV_DEVICE}"
+RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION DEVICE=${ENV_DEVICE} TF=${ENV_TF}"
 RUN --mount=type=cache,target=/root/.cache/uv \
     export UV_LINK_MODE=copy && \
-    bash /deps/src/dependencies/scripts/setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} DEVICE=${ENV_DEVICE}
+    bash /deps/src/dependencies/scripts/setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} DEVICE=${ENV_DEVICE} TF=${ENV_TF}
 
 # Now copy the remaining code (source files that may change frequently)
 COPY ${PACKAGE_DIR}/maxtext/ src/maxtext/
diff --git a/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile b/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
index a9c7ee82b9..40f040f57a 100644
--- a/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
+++ b/src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
@@ -44,6 +44,10 @@ ENV ENV_LIBTPU_VERSION=$LIBTPU_VERSION
 ARG DEVICE
 ENV ENV_DEVICE=$DEVICE
 
+# TODO: remove default once separate TF and TF-free nightly image workflows are established
+ARG TF=true
+ENV ENV_TF=$TF
+
 ARG PACKAGE_DIR
 ENV PACKAGE_DIR=$PACKAGE_DIR
 
@@ -65,10 +69,10 @@ COPY ${PACKAGE_DIR}/maxtext/integration/vllm/ src/maxtext/integration/vllm/
 COPY libtpu.so* /root/custom_libtpu/
 
 # Install dependencies - these steps are cached unless the copied files change
-RUN echo "Running command: bash setup.sh MODE=$ENV_MODE WORKFLOW=$ENV_WORKFLOW JAX_VERSION=$ENV_JAX_VERSION LIBTPU_VERSION=$ENV_LIBTPU_VERSION DEVICE=${ENV_DEVICE}"
+RUN echo "Running command: bash setup.sh MODE=$ENV_MODE WORKFLOW=$ENV_WORKFLOW JAX_VERSION=$ENV_JAX_VERSION LIBTPU_VERSION=$ENV_LIBTPU_VERSION DEVICE=${ENV_DEVICE} TF=${ENV_TF}"
 RUN --mount=type=cache,target=/root/.cache/uv \
     export UV_LINK_MODE=copy && \
-    bash /deps/src/dependencies/scripts/setup.sh MODE=${ENV_MODE} WORKFLOW=${ENV_WORKFLOW} JAX_VERSION=${ENV_JAX_VERSION} LIBTPU_VERSION=${ENV_LIBTPU_VERSION} DEVICE=${ENV_DEVICE}
+    bash /deps/src/dependencies/scripts/setup.sh MODE=${ENV_MODE} WORKFLOW=${ENV_WORKFLOW} JAX_VERSION=${ENV_JAX_VERSION} LIBTPU_VERSION=${ENV_LIBTPU_VERSION} DEVICE=${ENV_DEVICE} TF=${ENV_TF}
 
 # Now copy the remaining code (source files that may change frequently)
 COPY ${PACKAGE_DIR}/maxtext/ src/maxtext/
diff --git a/src/dependencies/extra_deps/cuda12_overrides.txt b/src/dependencies/extra_deps/cuda12_overrides.txt
index a9652c0c4d..bcfda7b4e0 100644
--- a/src/dependencies/extra_deps/cuda12_overrides.txt
+++ b/src/dependencies/extra_deps/cuda12_overrides.txt
@@ -2,7 +2,5 @@ datasets>=4.8.5
 fsspec==2026.2.0
 gcsfs==2026.2.0
 orbax-checkpoint>=0.12.1
-tensorflow==2.20.0
-tensorflow-text==2.20.1
 markdown>=3.10.2
 xprof<=2.23.0
diff --git a/src/dependencies/extra_deps/jetstream_github_deps.txt b/src/dependencies/extra_deps/jetstream_github_deps.txt
new file mode 100644
index 0000000000..bb2def4715
--- /dev/null
+++ b/src/dependencies/extra_deps/jetstream_github_deps.txt
@@ -0,0 +1 @@
+google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip
diff --git a/src/dependencies/extra_deps/pre_train_github_deps.txt b/src/dependencies/extra_deps/pre_train_github_deps.txt
index 676f2e58e7..813cbae48a 100644
--- a/src/dependencies/extra_deps/pre_train_github_deps.txt
+++ b/src/dependencies/extra_deps/pre_train_github_deps.txt
@@ -1,2 +1 @@
-google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip
 mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip
diff --git a/src/dependencies/extra_deps/tf_requirements.txt b/src/dependencies/extra_deps/tf_requirements.txt
new file mode 100644
index 0000000000..82ad6f54af
--- /dev/null
+++ b/src/dependencies/extra_deps/tf_requirements.txt
@@ -0,0 +1,4 @@
+tensorflow==2.20.0
+tensorflow-datasets
+tensorflow-text==2.20.1
+seqio
diff --git a/src/dependencies/extra_deps/tpu_overrides.txt b/src/dependencies/extra_deps/tpu_overrides.txt
index f2fa0c56cf..15805e11e7 100644
--- a/src/dependencies/extra_deps/tpu_overrides.txt
+++ b/src/dependencies/extra_deps/tpu_overrides.txt
@@ -2,7 +2,5 @@ datasets>=4.8.5
 fsspec==2026.2.0
 gcsfs==2026.2.0
 orbax-checkpoint>=0.12.4
-tensorflow==2.20.0
-tensorflow-text==2.20.1
 markdown>=3.10.2
 xprof<=2.23.0
diff --git a/src/dependencies/requirements/base_requirements/requirements.txt b/src/dependencies/requirements/base_requirements/requirements.txt
index 4c8a9fcc44..9d785716bf 100644
--- a/src/dependencies/requirements/base_requirements/requirements.txt
+++ b/src/dependencies/requirements/base_requirements/requirements.txt
@@ -35,12 +35,8 @@ pylint
 pytest
 pytype
 sentencepiece
-seqio
 tensorboard-plugin-profile
 tensorboardx
-tensorflow-datasets
-tensorflow-text
-tensorflow
 tiktoken
 tokamax!=0.1.0
 transformers>=5.8.0
diff --git a/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt b/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt
index 9b1cbec2fb..aee4d16d21 100644
--- a/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt
+++ b/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt
@@ -28,7 +28,6 @@ chex>=0.1.92
 click>=8.5.0
 cloud-accelerator-diagnostics>=0.1.1
 cloudpickle>=3.1.2
-clu>=0.0.12
 colorama>=0.4.6
 contourpy>=1.3.3
 coverage>=7.16.0
@@ -40,11 +39,8 @@ defusedxml>=0.7.1
 dill>=0.4.1
 distlib>=0.4.3
 distro>=1.9.0
-dm-tree>=0.1.10
 docstring-parser>=0.18.0
 drjax>=0.2.1
-editdistance>=0.8.1
-einops>=0.8.2
 einshape>=1.0
 etils>=1.14.0
 evaluate>=0.2.0
@@ -60,11 +56,11 @@ fsspec>=2026.2.0
 gast>=0.7.0
 gcsfs>=2026.2.0
 google-api-core>=2.34.0
-google-api-python-client>=2.199.0
+google-api-python-client>=2.200.0
 google-auth>=2.57.0
 google-auth-httplib2>=0.4.2
 google-auth-oauthlib>=1.4.1
-google-cloud-aiplatform>=2.0.1
+google-cloud-aiplatform>=2.1.0
 google-cloud-appengine-logging>=1.10.0
 google-cloud-audit-log>=0.6.2
 google-cloud-bigquery>=3.44.0
@@ -76,7 +72,7 @@ google-cloud-resource-manager>=1.18.0
 google-cloud-storage>=3.13.1
 google-cloud-storage-control>=1.14.0
 google-crc32c>=1.8.0
-google-genai>=2.20.0
+google-genai>=2.21.0
 google-pasta>=0.2.0
 google-resumable-media>=2.10.2
 googleapis-common-protos>=1.75.2
@@ -115,7 +111,7 @@ kiwisolver>=1.5.0
 latex2sympy2-extended>=1.11.0
 libclang>=18.1.1
 libcst>=1.9.0
-lxml>=6.1.2
+lxml>=6.1.3
 markdown>=3.10.2
 markdown-it-py>=4.2.0
 markupsafe>=3.0.3
@@ -149,7 +145,7 @@ nvidia-cuda-nvcc-cu12>=12.9.86 ; sys_platform == 'linux'
 nvidia-cuda-nvrtc-cu12>=12.9.86 ; sys_platform == 'linux'
 nvidia-cuda-runtime-cu12>=12.9.79 ; sys_platform == 'linux'
 nvidia-cudnn-cu12>=9.24.0.43 ; sys_platform == 'linux'
-nvidia-cudnn-frontend>=1.27.0
+nvidia-cudnn-frontend>=1.28.0
 nvidia-cufft-cu12>=11.4.1.4 ; sys_platform == 'linux'
 nvidia-cusolver-cu12>=11.7.5.82 ; sys_platform == 'linux'
 nvidia-cusparse-cu12>=12.5.10.65 ; sys_platform == 'linux'
@@ -171,12 +167,11 @@ parameterized>=0.9.0
 pathspec>=1.1.1
 pathwaysutils>=0.1.11
 pillow>=12.3.0
-platformdirs>=4.11.5
+platformdirs>=4.11.7
 pluggy>=1.6.0
 portpicker>=1.6.0
 pre-commit>=4.6.2
 prometheus-client>=0.26.0
-promise>=2.3
 propcache>=0.5.2
 proto-plus>=1.28.4
 protobuf>=6.33.6
@@ -191,7 +186,6 @@ pydantic>=2.13.5
 pydantic-core>=2.46.5
 pydot>=4.0.1
 pyelftools>=0.33
-pyglove>=0.4.5
 pygments>=2.21.0
 pyink>=26.5.1
 pylint>=4.0.8
@@ -206,7 +200,7 @@ pytokens>=0.4.1
 pytype>=2024.10.11
 pyyaml>=6.0.3
 qwix>=0.1.8
-regex>=2026.8.31
+regex>=2026.9.3
 requests>=2.34.2
 requests-oauthlib>=2.0.0
 responses>=0.18.0
@@ -215,11 +209,9 @@ safetensors>=0.8.0
 scipy>=1.18.0
 scipy-stubs>=1.18.0.1
 sentencepiece>=0.2.2
-seqio>=0.0.20
 setuptools>=84.0.0
 shellingham>=1.5.4
 shortuuid>=1.0.13
-simple-parsing>=0.1.9
 simplejson>=4.1.2
 six>=1.17.0
 sniffio>=1.3.1
@@ -228,14 +220,8 @@ starlette>=1.6.0
 sympy>=1.12
 tabulate>=0.10.0
 tenacity>=9.1.4
-tensorboard>=2.20.0
-tensorboard-data-server>=0.7.0
 tensorboard-plugin-profile>=2.13.0
 tensorboardx>=2.6.5
-tensorflow>=2.20.0
-tensorflow-datasets>=4.9.10
-tensorflow-metadata>=1.21.0
-tensorflow-text>=2.20.1
 tensorstore>=0.1.85
 termcolor>=3.3.0
 tiktoken>=0.14.0
@@ -259,12 +245,11 @@ uritemplate>=4.2.0
 urllib3>=2.7.0
 uvicorn>=0.52.4
 uvloop>=0.22.1
-virtualenv>=21.7.7
+virtualenv>=21.7.8
 wadler-lindig>=0.1.7
 websockets>=16.1.1
 werkzeug>=3.1.8
 wheel>=0.48.0
-wrapt>=2.4.0
 xprof>=2.20.0
 xxhash>=4.0.1
 yarl>=1.24.5
diff --git a/src/dependencies/requirements/generated_requirements/decoupled-requirements.txt b/src/dependencies/requirements/generated_requirements/decoupled-requirements.txt
index 2417cbca94..3ce06798ac 100644
--- a/src/dependencies/requirements/generated_requirements/decoupled-requirements.txt
+++ b/src/dependencies/requirements/generated_requirements/decoupled-requirements.txt
@@ -31,7 +31,6 @@ cheroot>=11.1.2
 chex>=0.1.92
 click>=8.5.0
 cloudpickle>=3.1.2
-clu>=0.0.12
 colorama>=0.4.6
 contourpy>=1.3.3
 coverage>=7.16.0
@@ -43,11 +42,8 @@ defusedxml>=0.7.1
 dill>=0.4.1
 distlib>=0.4.3
 distro>=1.9.0
-dm-tree>=0.1.10
 docstring-parser>=0.18.0
 drjax>=0.2.1
-editdistance>=0.8.1
-einops>=0.8.2
 einshape>=1.0
 etils>=1.14.0
 evaluate>=0.2.0
@@ -94,7 +90,7 @@ kiwisolver>=1.5.0
 latex2sympy2-extended>=1.11.0
 libclang>=18.1.1
 libcst>=1.9.0
-lxml>=6.1.2
+lxml>=6.1.3
 markdown>=3.10.2
 markdown-it-py>=4.2.0
 markupsafe>=3.0.3
@@ -134,12 +130,11 @@ parameterized>=0.9.0
 pathspec>=1.1.1
 pathwaysutils>=0.1.11
 pillow>=12.3.0
-platformdirs>=4.11.5
+platformdirs>=4.11.7
 pluggy>=1.6.0
 portpicker>=1.6.0
 pre-commit>=4.6.2
 prometheus-client>=0.26.0
-promise>=2.3
 propcache>=0.5.2
 protobuf>=6.33.6
 psutil>=7.2.2
@@ -153,7 +148,6 @@ pydantic>=2.13.5
 pydantic-core>=2.46.5
 pydot>=4.0.1
 pyelftools>=0.33
-pyglove>=0.4.5
 pygments>=2.21.0
 pyink>=26.5.1
 pylint>=4.0.8
@@ -168,7 +162,7 @@ pytokens>=0.4.1
 pytype>=2024.10.11
 pyyaml>=6.0.3
 qwix>=0.1.8
-regex>=2026.8.31
+regex>=2026.9.3
 requests>=2.34.2
 requests-oauthlib>=2.0.0
 responses>=0.18.0
@@ -177,11 +171,9 @@ safetensors>=0.8.0
 scipy>=1.18.0
 scipy-stubs>=1.18.0.1
 sentencepiece>=0.2.2
-seqio>=0.0.20
 setuptools>=84.0.0
 shellingham>=1.5.4
 shortuuid>=1.0.13
-simple-parsing>=0.1.9
 simplejson>=4.1.2
 six>=1.17.0
 sniffio>=1.3.1
@@ -190,14 +182,8 @@ starlette>=1.6.0
 sympy>=1.12
 tabulate>=0.10.0
 tenacity>=9.1.4
-tensorboard>=2.20.0
-tensorboard-data-server>=0.7.0
 tensorboard-plugin-profile>=2.13.0
 tensorboardx>=2.6.5
-tensorflow>=2.20.0
-tensorflow-datasets>=4.9.10
-tensorflow-metadata>=1.21.0
-tensorflow-text>=2.20.1
 tensorstore>=0.1.85
 termcolor>=3.3.0
 tiktoken>=0.14.0
@@ -218,12 +204,11 @@ uritemplate>=4.2.0
 urllib3>=2.7.0
 uvicorn>=0.52.4
 uvloop>=0.22.1
-virtualenv>=21.7.7
+virtualenv>=21.7.8
 wadler-lindig>=0.1.7
 websockets>=16.1.1
 werkzeug>=3.1.8
 wheel>=0.48.0
-wrapt>=2.4.0
 xxhash>=4.0.1
 yarl>=1.24.5
 zipp>=4.1.0
diff --git a/src/dependencies/requirements/generated_requirements/tpu-requirements.txt b/src/dependencies/requirements/generated_requirements/tpu-requirements.txt
index fc9ff57b2c..5b9413f72b 100644
--- a/src/dependencies/requirements/generated_requirements/tpu-requirements.txt
+++ b/src/dependencies/requirements/generated_requirements/tpu-requirements.txt
@@ -28,7 +28,6 @@ chex>=0.1.92
 click>=8.5.0
 cloud-accelerator-diagnostics>=0.1.1
 cloudpickle>=3.1.2
-clu>=0.0.12
 colorama>=0.4.6
 contourpy>=1.3.3
 coverage>=7.16.0
@@ -40,11 +39,8 @@ defusedxml>=0.7.1
 dill>=0.4.1
 distlib>=0.4.3
 distro>=1.9.0
-dm-tree>=0.1.10
 docstring-parser>=0.18.0
 drjax>=0.2.1
-editdistance>=0.8.1
-einops>=0.8.2
 einshape>=1.0
 etils>=1.14.0
 evaluate>=0.2.0
@@ -60,11 +56,11 @@ fsspec>=2026.2.0
 gast>=0.7.0
 gcsfs>=2026.2.0
 google-api-core>=2.34.0
-google-api-python-client>=2.199.0
+google-api-python-client>=2.200.0
 google-auth>=2.57.0
 google-auth-httplib2>=0.4.2
 google-auth-oauthlib>=1.4.1
-google-cloud-aiplatform>=2.0.1
+google-cloud-aiplatform>=2.1.0
 google-cloud-appengine-logging>=1.10.0
 google-cloud-audit-log>=0.6.2
 google-cloud-bigquery>=3.44.0
@@ -76,7 +72,7 @@ google-cloud-resource-manager>=1.18.0
 google-cloud-storage>=3.13.1
 google-cloud-storage-control>=1.14.0
 google-crc32c>=1.8.0
-google-genai>=2.20.0
+google-genai>=2.21.0
 google-pasta>=0.2.0
 google-resumable-media>=2.10.2
 googleapis-common-protos>=1.75.2
@@ -113,7 +109,7 @@ latex2sympy2-extended>=1.11.0
 libclang>=18.1.1
 libcst>=1.9.0
 libtpu>=0.0.46 ; platform_machine == 'x86_64' and sys_platform == 'linux'
-lxml>=6.1.2
+lxml>=6.1.3
 markdown>=3.10.2
 markdown-it-py>=4.2.0
 markupsafe>=3.0.3
@@ -155,12 +151,11 @@ parameterized>=0.9.0
 pathspec>=1.1.1
 pathwaysutils>=0.1.11
 pillow>=12.3.0
-platformdirs>=4.11.5
+platformdirs>=4.11.7
 pluggy>=1.6.0
 portpicker>=1.6.0
 pre-commit>=4.6.2
 prometheus-client>=0.26.0
-promise>=2.3
 propcache>=0.5.2
 proto-plus>=1.28.4
 protobuf>=6.33.6
@@ -175,7 +170,6 @@ pydantic>=2.13.5
 pydantic-core>=2.46.5
 pydot>=4.0.1
 pyelftools>=0.33
-pyglove>=0.4.5
 pygments>=2.21.0
 pyink>=26.5.1
 pylint>=4.0.8
@@ -190,7 +184,7 @@ pytokens>=0.4.1
 pytype>=2024.10.11
 pyyaml>=6.0.3
 qwix>=0.1.8
-regex>=2026.8.31
+regex>=2026.9.3
 requests>=2.34.2
 requests-oauthlib>=2.0.0
 responses>=0.18.0
@@ -199,11 +193,9 @@ safetensors>=0.8.0
 scipy>=1.18.0
 scipy-stubs>=1.18.0.1
 sentencepiece>=0.2.2
-seqio>=0.0.20
 setuptools>=84.0.0
 shellingham>=1.5.4
 shortuuid>=1.0.13
-simple-parsing>=0.1.9
 simplejson>=4.1.2
 six>=1.17.0
 sniffio>=1.3.1
@@ -212,14 +204,8 @@ starlette>=1.6.0
 sympy>=1.12
 tabulate>=0.10.0
 tenacity>=9.1.4
-tensorboard>=2.20.0
-tensorboard-data-server>=0.7.0
 tensorboard-plugin-profile>=2.13.0
 tensorboardx>=2.6.5
-tensorflow>=2.20.0
-tensorflow-datasets>=4.9.10
-tensorflow-metadata>=1.21.0
-tensorflow-text>=2.20.1
 tensorstore>=0.1.85
 termcolor>=3.3.0
 tiktoken>=0.14.0
@@ -240,12 +226,11 @@ uritemplate>=4.2.0
 urllib3>=2.7.0
 uvicorn>=0.52.4
 uvloop>=0.22.1
-virtualenv>=21.7.7
+virtualenv>=21.7.8
 wadler-lindig>=0.1.7
 websockets>=16.1.1
 werkzeug>=3.1.8
 wheel>=0.48.0
-wrapt>=2.4.0
 xprof>=2.20.0
 xxhash>=4.0.1
 yarl>=1.24.5
diff --git a/src/dependencies/scripts/docker_build_dependency_image.sh b/src/dependencies/scripts/docker_build_dependency_image.sh
index bc328f21d7..6a525efbd1 100644
--- a/src/dependencies/scripts/docker_build_dependency_image.sh
+++ b/src/dependencies/scripts/docker_build_dependency_image.sh
@@ -66,6 +66,9 @@ fi
 if [[ -z ${INCLUDE_TEST_ASSETS} ]]; then
   export INCLUDE_TEST_ASSETS=false
 fi
+if [[ -z ${TF} ]]; then
+  export TF=false
+fi
 
 # Create docker build arguments array
 docker_build_args=(
@@ -75,6 +78,7 @@ docker_build_args=(
   "JAX_VERSION=${JAX_VERSION}"
   "PACKAGE_DIR=${PACKAGE_DIR}"
   "INCLUDE_TEST_ASSETS=${INCLUDE_TEST_ASSETS}"
+  "TF=${TF}"
 )
 
 run_docker_build() {
diff --git a/src/dependencies/scripts/install_pre_train_extra_deps.py b/src/dependencies/scripts/install_pre_train_extra_deps.py
index 310e246b8b..0f40df6ae4 100644
--- a/src/dependencies/scripts/install_pre_train_extra_deps.py
+++ b/src/dependencies/scripts/install_pre_train_extra_deps.py
@@ -14,20 +14,31 @@
 
 """Installs extra dependencies from a requirements file using uv.
 
-This script is designed to install dependencies specified in 'dependencies/extra_deps/pre_train_*.txt'.
+This script is designed to install dependencies specified in 'dependencies/extra_deps/pre_train_*.txt'
+and optional TensorFlow dependencies from 'dependencies/extra_deps/tf_requirements.txt'.
 It first ensures 'uv' is installed and then uses it to install the packages listed in the requirements file.
 """
 
+import argparse
 import os
 import subprocess
 import sys
 
 
+def parse_args():
+  parser = argparse.ArgumentParser(description="Install pre-training extra dependencies.")
+  parser.add_argument(
+      "--with-tf",
+      action="store_true",
+      default=os.getenv("WITH_TF", "").lower() in ("true", "1", "yes"),
+      help="Install optional TensorFlow, TFDS, SeqIO, and JetStream dependencies.",
+  )
+  return parser.parse_args()
+
+
 def main():
-  """
-  Installs extra dependencies specified in 'dependencies/extra_deps/pre_train_*.txt' using uv.
-  It executes 'uv pip install -r <path_to_extra_deps.txt> --resolution=lowest'.
-  """
+  """Installs extra dependencies specified in 'dependencies/extra_deps/pre_train_*.txt' using uv."""
+  args = parse_args()
   current_dir = os.path.dirname(os.path.abspath(__file__))
   repo_root = os.path.abspath(os.path.join(current_dir, "..", ".."))
   github_deps_path = os.path.join(repo_root, "dependencies", "extra_deps", "pre_train_github_deps.txt")
@@ -56,20 +67,66 @@ def main():
 
   try:
     print(f"Installing Github dependencies: {' '.join(github_deps_command)}")
-    _ = subprocess.run(github_deps_command, check=True, capture_output=True, text=True)
+    subprocess.run(github_deps_command, check=True)
     print("Github dependencies installed successfully!")
   except subprocess.CalledProcessError as e:
     print("Failed to install extra dependencies.")
     print(f"Command '{' '.join(e.cmd)}' returned non-zero exit status {e.returncode}.")
-    print("--- Stderr ---")
-    print(e.stderr)
-    print("--- Stdout ---")
-    print(e.stdout)
     sys.exit(e.returncode)
   except (OSError, FileNotFoundError) as e:
     print(f"An OS-level error occurred while trying to run uv: {e}")
     sys.exit(1)
 
+  if args.with_tf:
+    tf_deps_path = os.path.join(repo_root, "dependencies", "extra_deps", "tf_requirements.txt")
+    if not os.path.exists(tf_deps_path):
+      raise FileNotFoundError(f"TensorFlow dependencies file not found at {tf_deps_path}")
+
+    tf_deps_command = [
+        sys.executable,
+        "-m",
+        "uv",
+        "pip",
+        "install",
+        "-r",
+        str(tf_deps_path),
+    ]
+    try:
+      print(f"Installing optional TensorFlow dependencies: {' '.join(tf_deps_command)}")
+      subprocess.run(tf_deps_command, check=True)
+      print("TensorFlow dependencies installed successfully!")
+    except subprocess.CalledProcessError as e:
+      print("Failed to install TensorFlow dependencies.")
+      print(f"Command '{' '.join(e.cmd)}' returned non-zero exit status {e.returncode}.")
+      sys.exit(e.returncode)
+    except (OSError, FileNotFoundError) as e:
+      print(f"An OS-level error occurred while trying to run uv: {e}")
+      sys.exit(1)
+
+    jetstream_deps_path = os.path.join(repo_root, "dependencies", "extra_deps", "jetstream_github_deps.txt")
+    if os.path.exists(jetstream_deps_path):
+      jetstream_deps_command = [
+          sys.executable,
+          "-m",
+          "uv",
+          "pip",
+          "install",
+          "-r",
+          str(jetstream_deps_path),
+          "--no-deps",
+      ]
+      try:
+        print(f"Installing JetStream dependencies: {' '.join(jetstream_deps_command)}")
+        subprocess.run(jetstream_deps_command, check=True)
+        print("JetStream dependencies installed successfully!")
+      except subprocess.CalledProcessError as e:
+        print("Failed to install JetStream dependencies.")
+        print(f"Command '{' '.join(e.cmd)}' returned non-zero exit status {e.returncode}.")
+        sys.exit(e.returncode)
+      except (OSError, FileNotFoundError) as e:
+        print(f"An OS-level error occurred while trying to run uv: {e}")
+        sys.exit(1)
+
 
 if __name__ == "__main__":
   main()
diff --git a/src/dependencies/scripts/setup.sh b/src/dependencies/scripts/setup.sh
index 9c0641edc8..ed0378263c 100644
--- a/src/dependencies/scripts/setup.sh
+++ b/src/dependencies/scripts/setup.sh
@@ -42,6 +42,9 @@
 # Install custom libtpu only
 ## bash src/dependencies/scripts/setup.sh MODE=libtpu-only LIBTPU_GCS_PATH=gs://my_custom_libtpu/libtpu.so
 
+# Install dependencies in dependencies/generated_requirements/tpu-requirements.txt + optional TensorFlow/JetStream
+## bash src/dependencies/scripts/setup.sh MODE=stable TF=true
+
 # ==================================
 # GPU EXAMPLES
 # ==================================
@@ -49,6 +52,9 @@
 # Install dependencies in dependencies/generated_requirements/cuda12-requirements.txt
 ## bash src/dependencies/scripts/setup.sh MODE=stable DEVICE=gpu
 
+# Install dependencies in dependencies/generated_requirements/cuda12-requirements.txt + optional TensorFlow/JetStream
+## bash src/dependencies/scripts/setup.sh MODE=stable DEVICE=gpu TF=true
+
 # Install dependencies in dependencies/generated_requirements/cuda12-requirements.txt + specified jax, jaxlib, jax-cuda12-plugin, jax-cuda12-pjrt
 ## bash src/dependencies/scripts/setup.sh MODE=stable DEVICE=gpu JAX_VERSION=0.4.13
 
@@ -159,6 +165,11 @@ if [[ -z "$WORKFLOW" ]]; then
   export WORKFLOW=pre-training
 fi
 
+# Set default value for $TF (default false for local setups; docker builds pass TF=true)
+if [[ -z "$TF" ]]; then
+  export TF="${tf:-false}"
+fi
+
 # Unset optional variables if set to NONE
 unset_optional_vars() {
     local optional_vars=("JAX_VERSION" "LIBTPU_VERSION" "LIBTPU_GCS_PATH")
@@ -216,7 +227,11 @@ install_maxtext_with_deps() {
     fi
     echo "Installing requirements from $dep_name"
     python3 -m uv pip install --resolution=lowest -r "$dep_name"
-    python3 -m src.dependencies.scripts.install_pre_train_extra_deps
+    if [[ "$TF" == "true" ]]; then
+        python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+    else
+        python3 -m src.dependencies.scripts.install_pre_train_extra_deps
+    fi
 
     install_maxtext_package_without_deps
 }
diff --git a/src/maxtext/common/gcloud_stub.py b/src/maxtext/common/gcloud_stub.py
index c87ba4123c..41f04e7798 100644
--- a/src/maxtext/common/gcloud_stub.py
+++ b/src/maxtext/common/gcloud_stub.py
@@ -148,8 +148,9 @@ def jetstream():
       if importlib.util.find_spec(mod) is None:
         if is_decoupled():
           print("[DECOUPLED NO-OP] jetstream: dependency missing; using stubs.")
-          return _jetstream_stubs()
-        raise ModuleNotFoundError(mod)
+        else:
+          print("[NO-OP] jetstream: dependency missing; using stubs.")
+        return _jetstream_stubs()
     from jetstream.core import config_lib  # type: ignore  # pylint: disable=import-outside-toplevel
     from jetstream.engine import engine_api, token_utils, tokenizer_api  # type: ignore  # pylint: disable=import-outside-toplevel
     from jetstream.engine.tokenizer_pb2 import TokenizerParameters, TokenizerType  # type: ignore  # pylint: disable=import-outside-toplevel
@@ -173,11 +174,12 @@ def jetstream():
     token_params_ns = SimpleNamespace(TokenizerParameters=TokenizerParameters, TokenizerType=TokenizerType)
     setattr(token_params_ns, "_IS_STUB", False)
     return config_lib, engine_api, token_utils, tokenizer_api, token_params_ns
-  except ModuleNotFoundError:
+  except (ModuleNotFoundError, ImportError) as e:
     if is_decoupled():
       print("[DECOUPLED NO-OP] jetstream: dependency missing; using stubs.")
-      return _jetstream_stubs()
-    raise
+    else:
+      print(f"[NO-OP] jetstream: dependency missing ({e}); using stubs.")
+    return _jetstream_stubs()
 
 
 # ---------------- GCS -----------------
diff --git a/src/maxtext/inference/decode.py b/src/maxtext/inference/decode.py
index 3034f5521e..9ccc549a6c 100644
--- a/src/maxtext/inference/decode.py
+++ b/src/maxtext/inference/decode.py
@@ -24,7 +24,7 @@
 
 from maxtext.configs import pyconfig
 from maxtext.common import profiler
-from maxtext.common.gcloud_stub import jetstream, is_decoupled
+from maxtext.common.gcloud_stub import jetstream
 from maxtext.inference.maxengine import maxengine
 from maxtext.multimodal import processor as mm_processor
 from maxtext.multimodal import utils as mm_utils
@@ -86,6 +86,14 @@ def _all_equals(elements: Sequence[jax.Array], target: jax.Array):
 
 
 def main(argv: Sequence[str]) -> None:
+  token_params_is_stub = getattr(_token_params_ns, "_IS_STUB", False)
+  engine_api_is_stub = getattr(engine_api, "_IS_STUB", False)
+  if token_params_is_stub or engine_api_is_stub:
+    raise RuntimeError(
+        "JetStream is not installed or disabled by DECOUPLE_GCLOUD=TRUE; maxtext.inference.decode requires JetStream. "
+        "Unset DECOUPLE_GCLOUD or install optional dependencies with: install_tpu_pre_train_extra_deps --with-tf"
+    )
+
   jax.config.update("jax_default_prng_impl", "unsafe_rbg")
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
 
@@ -119,13 +127,6 @@ def main(argv: Sequence[str]) -> None:
 
   metadata = engine.get_tokenizer()
   tokenizer_model = engine.build_tokenizer(metadata)
-  token_params_is_stub = getattr(_token_params_ns, "_IS_STUB", False)
-  engine_api_is_stub = getattr(engine_api, "_IS_STUB", False)
-  if is_decoupled() and (token_params_is_stub or engine_api_is_stub):
-    raise RuntimeError(
-        "JetStream disabled by DECOUPLE_GCLOUD=TRUE or stubbed; decode requires the JetStream tokenizer. "
-        "Unset DECOUPLE_GCLOUD or install JetStream to run decode."
-    )
 
   try:
     # TODO: update jetstream.engine.tokenizer_api.Tokenizer to maintain tokenizer state.
diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py
index 62c5fe0345..94d019ccc5 100644
--- a/src/maxtext/inference/maxengine/maxengine.py
+++ b/src/maxtext/inference/maxengine/maxengine.py
@@ -1885,10 +1885,10 @@ def get_tokenizer(self) -> Any:
     """
     token_params_is_stub = getattr(_token_params_ns, "_IS_STUB", False)
     engine_api_is_stub = getattr(engine_api, "_IS_STUB", False)
-    if is_decoupled() and (token_params_is_stub or engine_api_is_stub):
+    if token_params_is_stub or engine_api_is_stub:
       raise RuntimeError(
-          "JetStream disabled by DECOUPLE_GCLOUD=TRUE or stubbed; get_tokenizer is unsupported. "
-          "Unset DECOUPLE_GCLOUD or install JetStream to enable tokenizer functionality."
+          "JetStream is not installed or stubbed; get_tokenizer is unsupported. "
+          "Install optional dependencies with: install_tpu_pre_train_extra_deps --with-tf"
       )
     try:
       # pyrefly: ignore[missing-attribute]
@@ -1909,10 +1909,10 @@ def build_tokenizer(self, metadata: Any):  # return type depends on JetStream
     """Return a tokenizer"""
     token_params_is_stub = getattr(_token_params_ns, "_IS_STUB", False)
     engine_api_is_stub = getattr(engine_api, "_IS_STUB", False)
-    if is_decoupled() and (token_params_is_stub or engine_api_is_stub):
+    if token_params_is_stub or engine_api_is_stub:
       raise RuntimeError(
-          "JetStream disabled by DECOUPLE_GCLOUD=TRUE or stubbed; build_tokenizer is unsupported. "
-          "Unset DECOUPLE_GCLOUD or install JetStream to enable tokenizer functionality."
+          "JetStream is not installed or stubbed; build_tokenizer is unsupported. "
+          "Install optional dependencies with: install_tpu_pre_train_extra_deps --with-tf"
       )
     if metadata.tokenizer_type == TokenizerType.tiktoken:  # pyrefly: ignore[missing-attribute]
       return token_utils.TikToken(metadata)
diff --git a/tests/end_to_end/gpu/a3/test_llama2_7b.sh b/tests/end_to_end/gpu/a3/test_llama2_7b.sh
index c941cb8954..4b0fd5f554 100644
--- a/tests/end_to_end/gpu/a3/test_llama2_7b.sh
+++ b/tests/end_to_end/gpu/a3/test_llama2_7b.sh
@@ -10,6 +10,11 @@
 
 
 set -ex
+
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 idx=$(date +%Y-%m-%d-%H-%M)
 
 # Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to a GCS bucket that you own, this bucket will store all the files generated by MaxText during a run
diff --git a/tests/end_to_end/test_generate_param_only_checkpoint.sh b/tests/end_to_end/test_generate_param_only_checkpoint.sh
index 3b62cada13..c47cc63025 100644
--- a/tests/end_to_end/test_generate_param_only_checkpoint.sh
+++ b/tests/end_to_end/test_generate_param_only_checkpoint.sh
@@ -2,6 +2,11 @@
 
 set -uex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 helpFunction()
 {
   echo ""
diff --git a/tests/end_to_end/tpu/deepseek/v2-16b/test_deepseek.sh b/tests/end_to_end/tpu/deepseek/v2-16b/test_deepseek.sh
index c58310b487..3229b365e0 100644
--- a/tests/end_to_end/tpu/deepseek/v2-16b/test_deepseek.sh
+++ b/tests/end_to_end/tpu/deepseek/v2-16b/test_deepseek.sh
@@ -14,6 +14,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 export MODEL_NAME='deepseek2-16b'
 export TOKENIZER_PATH='deepseek-ai/DeepSeek-V2-Lite'
 
diff --git a/tests/end_to_end/tpu/deepseek/v3-671b/2_test_deepseek.sh b/tests/end_to_end/tpu/deepseek/v3-671b/2_test_deepseek.sh
index 718064d344..6b8eff111e 100644
--- a/tests/end_to_end/tpu/deepseek/v3-671b/2_test_deepseek.sh
+++ b/tests/end_to_end/tpu/deepseek/v3-671b/2_test_deepseek.sh
@@ -12,6 +12,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 export MODEL_NAME='deepseek3-671b'
 export TOKENIZER_PATH='deepseek-ai/DeepSeek-V3'
 
diff --git a/tests/end_to_end/tpu/deepseek/v3.2-671b/2_test_deepseek.sh b/tests/end_to_end/tpu/deepseek/v3.2-671b/2_test_deepseek.sh
index 3bceff1860..f597d4723c 100644
--- a/tests/end_to_end/tpu/deepseek/v3.2-671b/2_test_deepseek.sh
+++ b/tests/end_to_end/tpu/deepseek/v3.2-671b/2_test_deepseek.sh
@@ -9,6 +9,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 export MODEL_NAME='deepseek3.2-671b'
 export TOKENIZER_PATH='deepseek-ai/DeepSeek-V3.2'
 
diff --git a/tests/end_to_end/tpu/gemma/2b/test_gemma.sh b/tests/end_to_end/tpu/gemma/2b/test_gemma.sh
index ea4d2d783f..5998f4abeb 100644
--- a/tests/end_to_end/tpu/gemma/2b/test_gemma.sh
+++ b/tests/end_to_end/tpu/gemma/2b/test_gemma.sh
@@ -11,6 +11,11 @@
 
 
 set -ex
+
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 idx=$(date +%Y-%m-%d-%H-%M)
 export MODEL_VARIATION='2b'
 
diff --git a/tests/end_to_end/tpu/gemma/7b/2_test_gemma.sh b/tests/end_to_end/tpu/gemma/7b/2_test_gemma.sh
index 513212f309..9eb3b764e3 100644
--- a/tests/end_to_end/tpu/gemma/7b/2_test_gemma.sh
+++ b/tests/end_to_end/tpu/gemma/7b/2_test_gemma.sh
@@ -15,6 +15,11 @@
 # the subfolders names aka RUN_NAMEs are static. Please remember to change BASE_OUTPUT_PATH across different runs.
 
 set -ex
+
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 export MODEL_VARIATION='7b'
 
 if [ -z "${BASE_OUTPUT_PATH}" ]; then
diff --git a/tests/end_to_end/tpu/gemma2/27b/2_test_gemma.sh b/tests/end_to_end/tpu/gemma2/27b/2_test_gemma.sh
index 8f797144d0..4f8a28f189 100644
--- a/tests/end_to_end/tpu/gemma2/27b/2_test_gemma.sh
+++ b/tests/end_to_end/tpu/gemma2/27b/2_test_gemma.sh
@@ -15,6 +15,11 @@
 # the subfolders names aka RUN_NAMEs are static. Please remember to change BASE_OUTPUT_PATH across different runs.
 
 set -ex
+
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 export MODEL_VARIATION='27b'
 
 if [ -z "${BASE_OUTPUT_PATH}" ]; then
diff --git a/tests/end_to_end/tpu/gemma2/2b/test_gemma2.sh b/tests/end_to_end/tpu/gemma2/2b/test_gemma2.sh
index 092ed70333..1aba0e4c4f 100644
--- a/tests/end_to_end/tpu/gemma2/2b/test_gemma2.sh
+++ b/tests/end_to_end/tpu/gemma2/2b/test_gemma2.sh
@@ -10,6 +10,11 @@
 
 
 set -ex
+
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 idx=$(date +%Y-%m-%d-%H-%M)
 export MODEL_VARIATION='2b'
 
diff --git a/tests/end_to_end/tpu/gemma2/2b/test_gemma2_to_mt.sh b/tests/end_to_end/tpu/gemma2/2b/test_gemma2_to_mt.sh
index f3074b0de8..6613cda039 100644
--- a/tests/end_to_end/tpu/gemma2/2b/test_gemma2_to_mt.sh
+++ b/tests/end_to_end/tpu/gemma2/2b/test_gemma2_to_mt.sh
@@ -14,6 +14,11 @@
 
 
 set -ex
+
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 idx=$(date +%Y-%m-%d-%H-%M)
 MODEL_NAME='gemma2-2b'
 export MODEL_VARIATION='2b'
diff --git a/tests/end_to_end/tpu/gemma2/9b/2_test_gemma.sh b/tests/end_to_end/tpu/gemma2/9b/2_test_gemma.sh
index 7ba0868c91..9c3d4c873e 100644
--- a/tests/end_to_end/tpu/gemma2/9b/2_test_gemma.sh
+++ b/tests/end_to_end/tpu/gemma2/9b/2_test_gemma.sh
@@ -15,6 +15,11 @@
 # the subfolders names aka RUN_NAMEs are static. Please remember to change BASE_OUTPUT_PATH across different runs.
 
 set -ex
+
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 export MODEL_VARIATION='9b'
 
 # Installing torch for deps in forward_pass_logit_checker.py
diff --git a/tests/end_to_end/tpu/gemma3/12b/test_gemma3.sh b/tests/end_to_end/tpu/gemma3/12b/test_gemma3.sh
index 5276caa609..1ed8205a6e 100644
--- a/tests/end_to_end/tpu/gemma3/12b/test_gemma3.sh
+++ b/tests/end_to_end/tpu/gemma3/12b/test_gemma3.sh
@@ -9,6 +9,11 @@
 
 
 set -ex
+
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 idx=$(date +%Y-%m-%d-%H-%M)
 export MODEL_VARIATION='12b'
 export MODEL_NAME=gemma3-${MODEL_VARIATION}
diff --git a/tests/end_to_end/tpu/gemma3/27b/test_gemma3.sh b/tests/end_to_end/tpu/gemma3/27b/test_gemma3.sh
index 3ffd648752..33cd7bb6a0 100644
--- a/tests/end_to_end/tpu/gemma3/27b/test_gemma3.sh
+++ b/tests/end_to_end/tpu/gemma3/27b/test_gemma3.sh
@@ -9,6 +9,11 @@
 
 
 set -ex
+
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 idx=$(date +%Y-%m-%d-%H-%M)
 export MODEL_VARIATION='27b'
 export MODEL_NAME=gemma3-${MODEL_VARIATION}
diff --git a/tests/end_to_end/tpu/gemma3/4b/test_gemma3.sh b/tests/end_to_end/tpu/gemma3/4b/test_gemma3.sh
index 7efdf1a69b..146428e6d2 100644
--- a/tests/end_to_end/tpu/gemma3/4b/test_gemma3.sh
+++ b/tests/end_to_end/tpu/gemma3/4b/test_gemma3.sh
@@ -16,6 +16,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
 MODEL_NAME='gemma3-4b'
 
diff --git a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_multimodal_sft.sh b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_multimodal_sft.sh
index 72ca98c8ea..d582771f56 100644
--- a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_multimodal_sft.sh
+++ b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_multimodal_sft.sh
@@ -28,8 +28,10 @@ MULTIMODAL_SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/scanned_multimo
 # Non-Googlers please remember to point `DATASET_PATH` to the GCS bucket where you have your training data
 export DATASET_PATH=${DATASET_PATH:-gs://maxtext-dataset}
 
-# Step 1: Install google-jetstream
-python3 -m pip install google-jetstream@https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip --no-deps
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 
 # Step 2: Run inference on the original checkpoint converted from Hugging Face
 python3 -m maxtext.inference.decode \
diff --git a/tests/end_to_end/tpu/gemma4/26b/test_gemma4.sh b/tests/end_to_end/tpu/gemma4/26b/test_gemma4.sh
index 4cb7b2784e..6989a0edbb 100644
--- a/tests/end_to_end/tpu/gemma4/26b/test_gemma4.sh
+++ b/tests/end_to_end/tpu/gemma4/26b/test_gemma4.sh
@@ -16,6 +16,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
 MODEL_NAME='gemma4-26b'
 
diff --git a/tests/end_to_end/tpu/gpt_oss/120b/test_gpt_oss.sh b/tests/end_to_end/tpu/gpt_oss/120b/test_gpt_oss.sh
index 7181e3ac9a..4be8cc4771 100644
--- a/tests/end_to_end/tpu/gpt_oss/120b/test_gpt_oss.sh
+++ b/tests/end_to_end/tpu/gpt_oss/120b/test_gpt_oss.sh
@@ -14,6 +14,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 export MODEL_NAME='gpt-oss-120b'
 export TOKENIZER_PATH='openai/gpt-oss-120b'
 
diff --git a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh
index c143260282..affce9de99 100644
--- a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh
+++ b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh
@@ -15,6 +15,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 if [ -z "$1" ]; then
   echo "Error: run_id argument is required."
   exit 1
diff --git a/tests/end_to_end/tpu/llama2/13b/2_test_llama2_13b.sh b/tests/end_to_end/tpu/llama2/13b/2_test_llama2_13b.sh
index 0a1f363b4a..d02bcd3c1c 100644
--- a/tests/end_to_end/tpu/llama2/13b/2_test_llama2_13b.sh
+++ b/tests/end_to_end/tpu/llama2/13b/2_test_llama2_13b.sh
@@ -15,6 +15,11 @@
 # the subfolders names aka RUN_NAMEs are static. Please remember to change BASE_OUTPUT_PATH across different runs.
 
 set -ex
+
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 export MODEL_VARIATION='llama2-13b'
 
 if [ -z "${BASE_OUTPUT_PATH}" ]; then
diff --git a/tests/end_to_end/tpu/llama2/70b/2_test_llama2_70b.sh b/tests/end_to_end/tpu/llama2/70b/2_test_llama2_70b.sh
index e9d4e7daea..5a3e936c8c 100644
--- a/tests/end_to_end/tpu/llama2/70b/2_test_llama2_70b.sh
+++ b/tests/end_to_end/tpu/llama2/70b/2_test_llama2_70b.sh
@@ -16,6 +16,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 # Installing torch for deps in forward_pass_logit_checker.py
 python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
 
diff --git a/tests/end_to_end/tpu/llama2/7b/test_llama2_7b.sh b/tests/end_to_end/tpu/llama2/7b/test_llama2_7b.sh
index bd8456e2f5..826ed8bf53 100644
--- a/tests/end_to_end/tpu/llama2/7b/test_llama2_7b.sh
+++ b/tests/end_to_end/tpu/llama2/7b/test_llama2_7b.sh
@@ -11,6 +11,11 @@
 
 
 set -ex
+
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 idx=$(date +%Y-%m-%d-%H-%M)
 
 # Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to a GCS bucket that you own, this bucket will store all the files generated by MaxText during a run
diff --git a/tests/end_to_end/tpu/llama3.1/405b/2_test_llama3.1_405b.sh b/tests/end_to_end/tpu/llama3.1/405b/2_test_llama3.1_405b.sh
index 29a20dff6e..01fd8acdea 100644
--- a/tests/end_to_end/tpu/llama3.1/405b/2_test_llama3.1_405b.sh
+++ b/tests/end_to_end/tpu/llama3.1/405b/2_test_llama3.1_405b.sh
@@ -16,6 +16,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 # Installing torch for deps in forward_pass_logit_checker.py
 python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
 
diff --git a/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b.sh b/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b.sh
index 23e0f9074a..6b62525015 100644
--- a/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b.sh
+++ b/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b.sh
@@ -16,6 +16,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
 MODEL_NAME='llama3.1-70b'
 
diff --git a/tests/end_to_end/tpu/llama3.1/8b/2_test_llama3.1_8b.sh b/tests/end_to_end/tpu/llama3.1/8b/2_test_llama3.1_8b.sh
index 1c497a84f1..24ecca57d9 100644
--- a/tests/end_to_end/tpu/llama3.1/8b/2_test_llama3.1_8b.sh
+++ b/tests/end_to_end/tpu/llama3.1/8b/2_test_llama3.1_8b.sh
@@ -17,6 +17,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 # Installing torch for deps in forward_pass_logit_checker.py
 python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
 
diff --git a/tests/end_to_end/tpu/llama3.1/8b/3_test_llama3.1_8b.sh b/tests/end_to_end/tpu/llama3.1/8b/3_test_llama3.1_8b.sh
index 253a8b3b6f..7a1a21c042 100644
--- a/tests/end_to_end/tpu/llama3.1/8b/3_test_llama3.1_8b.sh
+++ b/tests/end_to_end/tpu/llama3.1/8b/3_test_llama3.1_8b.sh
@@ -3,6 +3,11 @@
 # huggingface-cli download meta-llama/Llama-3.1-8B-Instruct --local-dir $CHECKPOINT_ORIGINAL
 
 # Or download the DeepSeek llama 8B model
+
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 export CHECKPOINT_ORIGINAL=/mnt/disks/persist/checkpoints/huggingface/DeepSeek-R1-Distill-Llama-8B
 huggingface-cli download deepseek-ai/DeepSeek-R1-Distill-Llama-8B --local-dir $CHECKPOINT_ORIGINAL
 
diff --git a/tests/end_to_end/tpu/llama3.3/70b/2_test_llama3.3_70b.sh b/tests/end_to_end/tpu/llama3.3/70b/2_test_llama3.3_70b.sh
index ccd5267ad7..962f60ff87 100644
--- a/tests/end_to_end/tpu/llama3.3/70b/2_test_llama3.3_70b.sh
+++ b/tests/end_to_end/tpu/llama3.3/70b/2_test_llama3.3_70b.sh
@@ -16,6 +16,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 # Installing torch for deps in forward_pass_logit_checker.py
 python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
 
diff --git a/tests/end_to_end/tpu/llama3/70b/2_test_llama3_70b.sh b/tests/end_to_end/tpu/llama3/70b/2_test_llama3_70b.sh
index 41a3c8106a..acb21edf09 100644
--- a/tests/end_to_end/tpu/llama3/70b/2_test_llama3_70b.sh
+++ b/tests/end_to_end/tpu/llama3/70b/2_test_llama3_70b.sh
@@ -16,6 +16,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 # Installing torch for deps in forward_pass_logit_checker.py
 python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
 
diff --git a/tests/end_to_end/tpu/llama3/8b/2_test_llama3_8b.sh b/tests/end_to_end/tpu/llama3/8b/2_test_llama3_8b.sh
index 85c0bf18da..eb0b7297c0 100644
--- a/tests/end_to_end/tpu/llama3/8b/2_test_llama3_8b.sh
+++ b/tests/end_to_end/tpu/llama3/8b/2_test_llama3_8b.sh
@@ -16,6 +16,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 # Installing torch for deps in forward_pass_logit_checker.py
 python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
 
diff --git a/tests/end_to_end/tpu/mistral/7b/test_mistral-7b.sh b/tests/end_to_end/tpu/mistral/7b/test_mistral-7b.sh
index 29dee0065e..d9bcf4868d 100644
--- a/tests/end_to_end/tpu/mistral/7b/test_mistral-7b.sh
+++ b/tests/end_to_end/tpu/mistral/7b/test_mistral-7b.sh
@@ -10,6 +10,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 # Installing torch for deps in forward_pass_logit_checker.py
 python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
 
diff --git a/tests/end_to_end/tpu/mixtral/8x7b/2_test_mixtral.sh b/tests/end_to_end/tpu/mixtral/8x7b/2_test_mixtral.sh
index 6bd9862848..751db46ac9 100644
--- a/tests/end_to_end/tpu/mixtral/8x7b/2_test_mixtral.sh
+++ b/tests/end_to_end/tpu/mixtral/8x7b/2_test_mixtral.sh
@@ -12,6 +12,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 # Installing torch for deps in forward_pass_logit_checker.py
 python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
 
diff --git a/tests/end_to_end/tpu/qwen/dense/qwen2.5-14b/test_qwen2.5-14b.sh b/tests/end_to_end/tpu/qwen/dense/qwen2.5-14b/test_qwen2.5-14b.sh
index e0dd589a17..8cd8d63c9d 100644
--- a/tests/end_to_end/tpu/qwen/dense/qwen2.5-14b/test_qwen2.5-14b.sh
+++ b/tests/end_to_end/tpu/qwen/dense/qwen2.5-14b/test_qwen2.5-14b.sh
@@ -12,6 +12,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 export MODEL_NAME='qwen2.5-14b'
 export HF_MODEL_ID='Qwen/Qwen2.5-14B-Instruct'
 export TOKENIZER_PATH=${HF_MODEL_ID}
diff --git a/tests/end_to_end/tpu/qwen/dense/qwen2.5-7b/test_qwen2.5-7b.sh b/tests/end_to_end/tpu/qwen/dense/qwen2.5-7b/test_qwen2.5-7b.sh
index 7f4e619530..f8ce52e7b1 100644
--- a/tests/end_to_end/tpu/qwen/dense/qwen2.5-7b/test_qwen2.5-7b.sh
+++ b/tests/end_to_end/tpu/qwen/dense/qwen2.5-7b/test_qwen2.5-7b.sh
@@ -12,6 +12,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 export MODEL_NAME='qwen2.5-7b'
 export HF_MODEL_ID='Qwen/Qwen2.5-7B-Instruct'
 export TOKENIZER_PATH=${HF_MODEL_ID}
diff --git a/tests/end_to_end/tpu/qwen/moe/qwen3-omni-30b-a3b/2_test_qwen3_omni_30b_a3b.sh b/tests/end_to_end/tpu/qwen/moe/qwen3-omni-30b-a3b/2_test_qwen3_omni_30b_a3b.sh
index 0f902feccc..afd04334eb 100755
--- a/tests/end_to_end/tpu/qwen/moe/qwen3-omni-30b-a3b/2_test_qwen3_omni_30b_a3b.sh
+++ b/tests/end_to_end/tpu/qwen/moe/qwen3-omni-30b-a3b/2_test_qwen3_omni_30b_a3b.sh
@@ -14,6 +14,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 export MODEL_NAME="${MODEL_NAME:-qwen3-omni-30b-a3b}"
 export TOKENIZER_PATH="${TOKENIZER_PATH:-Qwen/Qwen3-Omni-30B-A3B-Instruct}"
 
diff --git a/tests/end_to_end/tpu/qwen/moe/qwen3.5-35b-a3b/2_test_qwen3.5_35b_a3b.sh b/tests/end_to_end/tpu/qwen/moe/qwen3.5-35b-a3b/2_test_qwen3.5_35b_a3b.sh
index 36629d7ee2..446d5a98eb 100644
--- a/tests/end_to_end/tpu/qwen/moe/qwen3.5-35b-a3b/2_test_qwen3.5_35b_a3b.sh
+++ b/tests/end_to_end/tpu/qwen/moe/qwen3.5-35b-a3b/2_test_qwen3.5_35b_a3b.sh
@@ -12,6 +12,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 export LIBTPU_INIT_ARGS="--xla_tpu_scoped_vmem_limit_kib=65536"
 export MODEL_NAME='qwen3.5-35b-a3b'
 export TOKENIZER_PATH='Qwen/Qwen3.5-35B-A3B'
diff --git a/tests/end_to_end/tpu/qwen/next/qwen3-next-80b-a3b/2_test_qwen3_next_80b_a3b.sh b/tests/end_to_end/tpu/qwen/next/qwen3-next-80b-a3b/2_test_qwen3_next_80b_a3b.sh
index 98529b16f1..95ff96214b 100644
--- a/tests/end_to_end/tpu/qwen/next/qwen3-next-80b-a3b/2_test_qwen3_next_80b_a3b.sh
+++ b/tests/end_to_end/tpu/qwen/next/qwen3-next-80b-a3b/2_test_qwen3_next_80b_a3b.sh
@@ -12,6 +12,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 export PYTHONPATH=$PYTHONPATH:$(pwd)/src
 
 export MODEL_NAME='qwen3-next-80b-a3b'
diff --git a/tests/end_to_end/tpu/qwen3/30b/test_qwen3.sh b/tests/end_to_end/tpu/qwen3/30b/test_qwen3.sh
index 26d18b6939..b6282b7748 100755
--- a/tests/end_to_end/tpu/qwen3/30b/test_qwen3.sh
+++ b/tests/end_to_end/tpu/qwen3/30b/test_qwen3.sh
@@ -16,6 +16,11 @@
 
 set -ex
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
 MODEL_NAME='qwen3-30b-a3b-base'
 
diff --git a/tests/end_to_end/tpu/qwen3/4b/test_qwen3_to_mt.sh b/tests/end_to_end/tpu/qwen3/4b/test_qwen3_to_mt.sh
index 8d5fc39e8b..ceb38bce37 100644
--- a/tests/end_to_end/tpu/qwen3/4b/test_qwen3_to_mt.sh
+++ b/tests/end_to_end/tpu/qwen3/4b/test_qwen3_to_mt.sh
@@ -13,6 +13,11 @@
 # export HF_TOKEN=<Hugging Face access token>
 
 set -ex
+
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 idx=$(date +%Y-%m-%d-%H-%M)
 MODEL_NAME='qwen3-4b'
 export MODEL_VARIATION='4b'
diff --git a/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_multimodal_sft.sh b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_multimodal_sft.sh
index 5595462b43..70c0cffb1f 100644
--- a/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_multimodal_sft.sh
+++ b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_multimodal_sft.sh
@@ -28,8 +28,10 @@ MULTIMODAL_UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/unscanned_mul
 # Non-Googlers please remember to point `DATASET_PATH` to the GCS bucket where you have your training data
 export DATASET_PATH=gs://maxtext-dataset
 
-# Step 1: Install google-jetstream
-python3 -m pip install google-jetstream@https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip --no-deps
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 
 # Step 2: Run inference on the original checkpoint converted from Hugging Face
 python3 -m maxtext.inference.decode \
diff --git a/tests/end_to_end/tpu/test_decode_save_quantized_ckpt.sh b/tests/end_to_end/tpu/test_decode_save_quantized_ckpt.sh
index ac3bdf3ff7..56173bca2b 100644
--- a/tests/end_to_end/tpu/test_decode_save_quantized_ckpt.sh
+++ b/tests/end_to_end/tpu/test_decode_save_quantized_ckpt.sh
@@ -2,6 +2,11 @@
 
 # Example run: bash tests/end_to_end/tpu/test_decode_save_quantized_ckpt.sh -m llama2-70b -r 070924 -n
 
+
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 dry_run=false
 model='llama2-7b'
 run_name="test_quant_ckpt"
diff --git a/tests/end_to_end/tpu/test_sft_trainer.sh b/tests/end_to_end/tpu/test_sft_trainer.sh
index 1b23a2959b..2d052203db 100755
--- a/tests/end_to_end/tpu/test_sft_trainer.sh
+++ b/tests/end_to_end/tpu/test_sft_trainer.sh
@@ -14,6 +14,11 @@
 
 set -xe
 
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
+
 RUN_NAME=sft-$(date +%Y-%m-%d-%H-%M-%S)
 PER_DEVICE_BATCH_SIZE=1
 LOSS_THRESHOLD=100.0 # Set to large value so test is guaranteed to pass
diff --git a/tests/inference/test_llama2_7b_bf16.sh b/tests/inference/test_llama2_7b_bf16.sh
index b02fa32d9a..a9071cf727 100755
--- a/tests/inference/test_llama2_7b_bf16.sh
+++ b/tests/inference/test_llama2_7b_bf16.sh
@@ -1,5 +1,10 @@
 #!/bin/bash
 
+
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 CONFIG_PATH="${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext}/configs/base.yml"
 if [ "${DECOUPLE_GCLOUD^^}" = "TRUE" ]; then
   CONFIG_PATH="${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext}/configs/decoupled_base_test.yml"
diff --git a/tests/inference/test_llama2_7b_int8.sh b/tests/inference/test_llama2_7b_int8.sh
index 8e11e6ab48..3ecf620618 100755
--- a/tests/inference/test_llama2_7b_int8.sh
+++ b/tests/inference/test_llama2_7b_int8.sh
@@ -1,5 +1,10 @@
 #!/bin/bash
 
+
+# Ensure JetStream and dependencies are installed for inference.decode
+if ! python3 -c "import jetstream" &>/dev/null; then
+    python3 -m src.dependencies.scripts.install_pre_train_extra_deps --with-tf
+fi
 CONFIG_PATH="${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext}/configs/base.yml"
 if [ "${DECOUPLE_GCLOUD^^}" = "TRUE" ]; then
   CONFIG_PATH="${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext}/configs/decoupled_base_test.yml"
diff --git a/tests/integration/decode_tests.py b/tests/integration/decode_tests.py
index 32c04fe155..d3db8bcc33 100644
--- a/tests/integration/decode_tests.py
+++ b/tests/integration/decode_tests.py
@@ -22,12 +22,17 @@
 from absl.testing import absltest
 from contextlib import redirect_stdout
 
+from maxtext.common.gcloud_stub import jetstream
 from maxtext.inference.decode import main as decode_main
 from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT
 from tests.utils.test_helpers import get_test_config_path, get_test_dataset_path, get_test_base_output_directory
 
 pytestmark = [pytest.mark.tpu_only, pytest.mark.external_serving, pytest.mark.integration_test]
 
+config_lib, _, _, _, _ = jetstream()
+if getattr(config_lib, "_IS_STUB", False):
+  pytest.skip("JetStream is not installed (stubbed)", allow_module_level=True)
+
 
 class DecodeTests(unittest.TestCase):
   """Tests decode with various configs."""
diff --git a/tests/integration/generate_param_only_checkpoint_test.py b/tests/integration/generate_param_only_checkpoint_test.py
index 6204144ce0..1d5ab9d6a7 100644
--- a/tests/integration/generate_param_only_checkpoint_test.py
+++ b/tests/integration/generate_param_only_checkpoint_test.py
@@ -21,7 +21,11 @@
 import pytest
 
 from maxtext.inference.decode import main as decode_main
-from maxtext.common.gcloud_stub import is_decoupled
+from maxtext.common.gcloud_stub import is_decoupled, jetstream
+
+config_lib, _, _, _, _ = jetstream()
+if getattr(config_lib, "_IS_STUB", False):
+  pytest.skip("JetStream is not installed (stubbed)", allow_module_level=True)
 from maxtext.trainers.pre_train.train import main as train_main
 from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT
 from maxtext.utils.generate_param_only_checkpoint import main as generate_param_only_ckpt_main
diff --git a/tests/integration/smoke/inference_microbenchmark_smoke_test.py b/tests/integration/smoke/inference_microbenchmark_smoke_test.py
index 6d039f8b70..2f0f46d6a2 100644
--- a/tests/integration/smoke/inference_microbenchmark_smoke_test.py
+++ b/tests/integration/smoke/inference_microbenchmark_smoke_test.py
@@ -21,10 +21,14 @@
 
 from maxtext.configs import pyconfig
 from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR, MAXTEXT_ASSETS_ROOT
-from maxtext.common.gcloud_stub import is_decoupled
+from maxtext.common.gcloud_stub import is_decoupled, jetstream
 
 pytestmark = [pytest.mark.external_serving]
 
+config_lib, _, _, _, _ = jetstream()
+if getattr(config_lib, "_IS_STUB", False):
+  pytest.skip("JetStream is not installed (stubbed)", allow_module_level=True)
+
 # Conditional import: only load when not in decoupled mode to avoid collection errors.
 # inference_microbenchmark depends on prefill_packing, which requires JetStream.
 if not is_decoupled():
diff --git a/tests/unit/offline_engine_test.py b/tests/unit/offline_engine_test.py
index 2d67d1fcec..474c83f5f5 100644
--- a/tests/unit/offline_engine_test.py
+++ b/tests/unit/offline_engine_test.py
@@ -22,10 +22,14 @@
 import jax.numpy as jnp
 import numpy as np
 from maxtext.configs import pyconfig
-from maxtext.common.gcloud_stub import is_decoupled
+from maxtext.common.gcloud_stub import is_decoupled, jetstream
 
 pytestmark = [pytest.mark.external_serving]
 
+config_lib, _, _, _, _ = jetstream()
+if getattr(config_lib, "_IS_STUB", False):
+  pytest.skip("JetStream is not installed (stubbed)", allow_module_level=True)
+
 # Conditional import: only load when not in decoupled mode to avoid collection errors.
 # offline_engine depends on prefill_packing, which requires JetStream.
 if not is_decoupled():