jax

4676543d - Update hypothesis test util to disable sharding after a failure (in shrinking phase) so that it can track down the failure correctly and not get confused by skipped tests.

Update hypothesis test util to disable sharding after a failure (in shrinking phase) so that it can track down the failure correctly and not get confused by skipped tests. I guess this could lead to timeouts, but only when the test already failed. PiperOrigin-RevId: 925458016