From fa68c1f882660285564d5d005ce06e4db51ea1cc Mon Sep 17 00:00:00 2001
From: Skye Wanderman-Milne <skyewm@google.com>
Date: Mon, 1 May 2023 13:30:49 -0700
Subject: [PATCH] Bump up lax_test TPU sharding to avoid asan timeouts

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

diff --git a/tests/BUILD b/tests/BUILD
index e75480316..9e7070e3c 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -419,7 +419,7 @@ jax_test(
     shard_count = {
         "cpu": 40,
         "gpu": 40,
-        "tpu": 30,
+        "tpu": 40,
         "iree": 40,
     },
     deps = [