diff --git a/.github/workflows/UploadDockerImages.yml b/.github/workflows/UploadDockerImages.yml
index 931ac4000..145de1300 100644
--- a/.github/workflows/UploadDockerImages.yml
+++ b/.github/workflows/UploadDockerImages.yml
@@ -102,6 +102,7 @@ jobs:
       dockerfile: ${{ matrix.dockerfile }}
       maxtext_sha: ${{ needs.setup.outputs.maxtext_sha }}
       image_date: ${{ needs.setup.outputs.image_date }}
+      base_image: gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:${{ needs.setup.outputs.image_date }}
       is_post_training: true
 
   gpu-pre-training:
diff --git a/.github/workflows/build_and_push_docker_image.yml b/.github/workflows/build_and_push_docker_image.yml
index 0168dccea..6c2fedc57 100644
--- a/.github/workflows/build_and_push_docker_image.yml
+++ b/.github/workflows/build_and_push_docker_image.yml
@@ -37,6 +37,10 @@ on:
       image_date:
         required: true
         type: string
+      base_image:
+        required: false
+        type: string
+        default: ''
       is_post_training:
         required: false
         type: boolean
@@ -114,7 +118,8 @@ jobs:
             MODE=${{ inputs.build_mode }}
             JAX_VERSION=NONE
             LIBTPU_VERSION=NONE
-            BASEIMAGE=gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:${{ inputs.image_date }}
+            INCLUDE_TEST_ASSETS=true
+            ${{ inputs.base_image != '' && format('BASEIMAGE={0}', inputs.base_image) || '' }}
 
       - name: Add tags to Docker image
         if: steps.check.outputs.should_run == 'true'
diff --git a/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile b/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
index a1abd24ac..5cb4368f6 100644
--- a/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
+++ b/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile
@@ -58,5 +58,14 @@ RUN --mount=type=cache,target=/root/.cache/pip bash /deps/tools/setup/setup.sh M
 # Now copy the remaining code (source files that may change frequently)
 COPY . .
 
+# Download test assets from GCS if building image with test assets
+ARG INCLUDE_TEST_ASSETS=false
+RUN if [ "$INCLUDE_TEST_ASSETS" = "true" ]; then \
+        echo "Downloading test assets from GCS..."; \
+        if ! gcloud storage cp -r gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT}"; then \
+        echo "WARNING: Failed to download test assets from GCS. These files are only used for end-to-end tests; you may not have access to the bucket."; \
+        fi; \
+    fi
+
 # Install (editable) MaxText
 RUN test -f '/tmp/venv_created' && "$(tail -n1 /tmp/venv_created)"/bin/activate ; pip install --no-dependencies -e .
diff --git a/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile b/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
index 1d6db625b..db2fe3824 100644
--- a/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
+++ b/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile
@@ -55,5 +55,14 @@ RUN --mount=type=cache,target=/root/.cache/pip bash /deps/tools/setup/setup.sh M
 # Now copy the remaining code (source files that may change frequently)
 COPY . .
 
+# Download test assets from GCS if building image with test assets
+ARG INCLUDE_TEST_ASSETS=false
+RUN if [ "$INCLUDE_TEST_ASSETS" = "true" ]; then \
+        echo "Downloading test assets from GCS..."; \
+        if ! gcloud storage cp -r gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT}"; then \
+        echo "WARNING: Failed to download test assets from GCS. These files are only used for end-to-end tests; you may not have access to the bucket."; \
+        fi; \
+    fi
+
 # Install (editable) MaxText
 RUN test -f '/tmp/venv_created' && "$(tail -n1 /tmp/venv_created)"/bin/activate ; pip install --no-dependencies -e .