From da4e79a6251ae71504061b496ea34f20dc06fa60 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Mon, 22 Aug 2022 06:51:21 -0700
Subject: [PATCH] Increase some test shardings to reduce CI timeouts under
 asan/tsan/msan.

PiperOrigin-RevId: 469173324
---
 tests/BUILD | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/BUILD b/tests/BUILD
index 88a8de6dc..d95eb8f12 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -317,9 +317,9 @@ jax_test(
     name = "lax_control_flow_test",
     srcs = ["lax_control_flow_test.py"],
     shard_count = {
-        "cpu": 10,
+        "cpu": 20,
         "gpu": 20,
-        "tpu": 10,
+        "tpu": 20,
         "iree": 10,
     },
 )
@@ -529,9 +529,9 @@ jax_test(
     name = "pmap_test",
     srcs = ["pmap_test.py"],
     shard_count = {
-        "cpu": 15,
+        "cpu": 30,
         "gpu": 30,
-        "tpu": 15,
+        "tpu": 30,
     },
     tags = ["multiaccelerator"],
     deps = [