From 6608242f95445414c225d8b6ab82dcc6898dcd3a Mon Sep 17 00:00:00 2001
From: Jake VanderPlas <vanderplas@google.com>
Date: Wed, 15 Feb 2023 14:57:54 -0800
Subject: [PATCH] sparse_test: reduce num_generated_cases to avoid timeouts

PiperOrigin-RevId: 509941080
---
 tests/BUILD | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tests/BUILD b/tests/BUILD
index 16ce56189..49475f902 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -755,6 +755,13 @@ jax_test(
         ],
         "tpu": ["optonly"],
     },
+    # Use fewer cases to prevent timeouts.
+    backend_variant_args = {
+        "cpu": ["--jax_num_generated_cases=40"],
+        "cpu_x32": ["--jax_num_generated_cases=40"],
+        "cpu_no_jax_array": ["--jax_num_generated_cases=40"],
+        "gpu": ["--jax_num_generated_cases=40"],
+    },
     shard_count = {
         "cpu": 50,
         "gpu": 50,