From 904cd4b98de0e5b47d8175ac097836c8d1f913d7 Mon Sep 17 00:00:00 2001
From: Adam Paszke <apaszke@google.com>
Date: Thu, 5 Jan 2023 04:12:49 -0800
Subject: [PATCH] Internal change

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

diff --git a/tests/BUILD b/tests/BUILD
index 489af3925..4b36d048a 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -731,8 +731,8 @@ jax_test(
     },
     enable_configs = ["cpu_jit_pjit_api_merge"],
     shard_count = {
-        "cpu": 40,
-        "gpu": 40,
+        "cpu": 50,
+        "gpu": 50,
         "tpu": 50,
         "iree": 10,
     },