From 117da44712500d89f6101c3a4109bdaf4eb06688 Mon Sep 17 00:00:00 2001
From: Adam Paszke <apaszke@google.com>
Date: Wed, 20 Jul 2022 04:30:50 -0700
Subject: [PATCH] Internal change

PiperOrigin-RevId: 462110048
---
 tests/BUILD | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/BUILD b/tests/BUILD
index 1f18042a0..7a65df699 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -457,7 +457,7 @@ jax_test(
         ],
     },
     shard_count = {
-        "cpu": 20,
+        "cpu": 25,
         "gpu": 40,
         "tpu": 10,
         "iree": 20,
@@ -683,6 +683,7 @@ jax_test(
     args = ["--jax_bcoo_cusparse_lowering=true"],
     shard_count = {
         "gpu": 20,
+        "tpu": 10,
     },
     deps = [
         "//jax:experimental_sparse",