Maxtext gpu tests on jsts

[Rohan-Bierneni]

Description

Maxtext GPU tests are built using the stable docker image. Since we have stable stack image now we can use that for a3 and a3 plus tests. However, first will verify that A3plus is supported, then if tests are green will follow up with A3 on stable stack.

Fixes: b/392001678

Tests

The tests in jax_stable_stack_gpu_e2e XLML DAG pass:
https://screenshot.googleplex.com/66JtkDem5zNNixj

Note: Recent failures in the DAG are due to bad pod phase

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run one-shot tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.