test te 2.6.0 jax

Author: Rohan-Bierneni

.github/workflows/UnitTests.yml

@@ -69,7 +69,7 @@ jobs:
  # IMPORTANT: Replace with the label for your runner (e.g., v5p-8)
  runs-on: ["linux-x86-a3-megagpu-h100-8gpu"]
  container:
- image: gcr.io/tpu-prod-env-multipod/maxdiffusion_stable_stack_candidate:jax0.7.2_rev3_gpu
+ image: gcr.io/tpu-prod-env-multipod/maxdiffusion_stable_stack_candidate:jax0.7.2_rev4_gpu
  steps:
  - name: Checkout Repository
  uses: actions/checkout@v4