From 5a7bedca376c6f7fe82071995bc3faed5e209d29 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Fri, 8 Jul 2022 08:29:56 -0700
Subject: [PATCH] Increase shard_count for sparse_test_gpu to 20.

https://github.com/google/jax/commit/1918d39765f6775de778625ecc4f197b2e104a6f updated the wrong test!

This test is close to the timeout in the GPU CI and flakes sometimes.

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

diff --git a/tests/BUILD b/tests/BUILD
index b0609eb43..6c990df86 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -375,7 +375,7 @@ jax_test(
     },
     shard_count = {
         "cpu": 10,
-        "gpu": 20,
+        "gpu": 10,
         "tpu": 10,
         "iree": 10,
     },
@@ -655,7 +655,7 @@ jax_test(
     },
     shard_count = {
         "cpu": 10,
-        "gpu": 10,
+        "gpu": 20,
         "tpu": 10,
         "iree": 10,
     },