update nightly image (GoogleCloudPlatform#820)

shuningjin · web-flow · commit df92d36ad455 · 2025-08-18T15:57:55.000-07:00
diff --git a/dags/sparsity_diffusion_devx/maxtext_moe_tpu_e2e.py b/dags/sparsity_diffusion_devx/maxtext_moe_tpu_e2e.py
@@ -50,7 +50,7 @@
   )
   docker_image = {
       "stable": DockerImage.MAXTEXT_TPU_JAX_STABLE_STACK.value,
-      "nightly": DockerImage.MAXTEXT_TPU_JAX_NIGHTLY.value,
+      "nightly": DockerImage.MAXTEXT_TPU_STABLE_STACK_NIGHTLY_JAX.value,
   }
 
   # Unchained tests

Original file line number	Diff line number	Diff line change
`@@ -50,7 +50,7 @@`
`50`	`50`	`)`
`51`	`51`	`docker_image = {`
`52`	`52`	`"stable": DockerImage.MAXTEXT_TPU_JAX_STABLE_STACK.value,`
`53`		`- "nightly": DockerImage.MAXTEXT_TPU_JAX_NIGHTLY.value,`
	`53`	`+ "nightly": DockerImage.MAXTEXT_TPU_STABLE_STACK_NIGHTLY_JAX.value,`
`54`	`54`	`}`
`55`	`55`
`56`	`56`	`# Unchained tests`