jax

e1f037bf - Reduce the max_examples for splash attention tests

Reduce the max_examples for splash attention tests They are too expensive to run all the time. We have 25 test methods, so it's ok to drop the number of examples each one of them covers. PiperOrigin-RevId: 707124466