Increase some test shardings to reduce CI timeouts under asan/tsan/msan.
Increase some test shardings to reduce CI timeouts under asan/tsan/msan.

PiperOrigin-RevId: 469173324

google/jax,google/jax,google/jax,google/jax