diff --git a/tests/BUILD b/tests/BUILD
index 27f2c115d..f8de3810d 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -704,6 +704,11 @@ jax_test(
             "notsan",  # Times out
         ],
     },
+    # Use fewer cases to prevent timeouts.
+    backend_variant_args = {
+        "cpu": ["--jax_num_generated_cases=40"],
+        "cpu_x32": ["--jax_num_generated_cases=40"],
+    },
     shard_count = {
         "cpu": 30,
         "gpu": 30,