From 8068e4638c924cfbe3dafb080b26c9024f9fea96 Mon Sep 17 00:00:00 2001
From: Sharad Vikram <sharadmv@google.com>
Date: Wed, 17 Aug 2022 10:45:47 -0700
Subject: [PATCH] Re-bump shard count for `pmap_test`

PiperOrigin-RevId: 468239588
---
 tests/BUILD | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/BUILD b/tests/BUILD
index 3fcdc0872..2716631c4 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -526,9 +526,9 @@ jax_test(
     name = "pmap_test",
     srcs = ["pmap_test.py"],
     shard_count = {
-        "cpu": 10,
-        "gpu": 10,
-        "tpu": 10,
+        "cpu": 15,
+        "gpu": 30,
+        "tpu": 15,
     },
     tags = ["multiaccelerator"],
     deps = [