From 864849d0752c4609e0b2d709b0b0bb25c96d6e7e Mon Sep 17 00:00:00 2001
From: Yash Katariya <yashkatariya@google.com>
Date: Sun, 11 Sep 2022 14:30:18 -0700
Subject: [PATCH] Increase the GPU shard count to 40 since its timing out every
 once in a while in kokoro.

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

diff --git a/tests/BUILD b/tests/BUILD
index e9b2988a9..1c052c6f2 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -318,7 +318,7 @@ jax_test(
     srcs = ["lax_control_flow_test.py"],
     shard_count = {
         "cpu": 30,
-        "gpu": 30,
+        "gpu": 40,
         "tpu": 30,
         "iree": 10,
     },