Increase `shard_count` to `4` for `jet_test` on GPU.
PiperOrigin-RevId: 971223193
B
Benjamin Chetioui committed
070e1edd4cbf88a41c783feb077c0898bd287575
Parent: 950509f
Committed by jax authors <google-ml-automation@google.com>
on 8/26/2026, 12:51:48 PM