From c18f7916a322d469378edd5467eeb6e99d9d8566 Mon Sep 17 00:00:00 2001
From: Roy Frostig <frostig@google.com>
Date: Tue, 30 Apr 2024 12:34:20 -0700
Subject: [PATCH] bump shard count for `random_lax_test`

PiperOrigin-RevId: 629495786
---
 tests/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/BUILD b/tests/BUILD
index 6ff21ab9f..86e61cc7f 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -810,7 +810,7 @@ jax_test(
     },
     shard_count = {
         "cpu": 40,
-        "gpu": 30,
+        "gpu": 40,
         "tpu": 40,
     },
     tags = ["noasan"],  # Times out