Use jax nightly until new jax version releases (GoogleCloudPlatform#960)

Author: Rohan-Bierneni

dags/sparsity_diffusion_devx/maxdiffusion_tpu_e2e.py

@@ -111,7 +111,7 @@
               f"output_dir={sdxl_base_output_dir}",
           ),
           test_name=sdxl_run_name_prefix,
-          docker_image=DockerImage.MAXDIFFUSION_TPU_JAX_STABLE_STACK.value,
+          docker_image=DockerImage.MAXDIFFUSION_TPU_STABLE_STACK_NIGHTLY_JAX.value,
           test_owner=test_owner.PARAM_B,
           tensorboard_summary_config=sdxl_tensorboard_summary_config,
       ).run_with_name_gen_and_quarantine(
@@ -133,7 +133,7 @@
               f"LOSS_THRESHOLD=100",
           ),
           test_name=sdxl_nan_run_name_prefix,
-          docker_image=DockerImage.MAXDIFFUSION_TPU_JAX_STABLE_STACK.value,
+          docker_image=DockerImage.MAXDIFFUSION_TPU_STABLE_STACK_NIGHTLY_JAX.value,
           test_owner=test_owner.PARAM_B,
           tensorboard_summary_config=sdxl_nan_tensorboard_summary_config,
       ).run_with_name_gen_and_quarantine(