random_test: Split into two so that each target is small enough to fi…

…t within a medium timeout. PiperOrigin-RevId: 571146766
google · Oct 5, 2023 · 8f911e1 · 8f911e1
1 parent 1512650
commit 8f911e1
Show file tree

Hide file tree

Showing 3 changed files with 1,412 additions and 1,310 deletions.
diff --git a/tests/BUILD b/tests/BUILD
@@ -704,10 +704,28 @@ 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,
+        "tpu": 40,
+        "iree": 30,
+    },
+    tags = ["noasan"],  # Times out
+)
+
+jax_test(
+    name = "random_lax_test",
+    srcs = ["random_lax_test.py"],
+    backend_tags = {
+        "cpu": [
+            "notsan",  # Times out
+            "nomsan",  # Times out
+        ],
+        "tpu": [
+            "optonly",
+            "nomsan",  # Times out
+            "notsan",  # Times out
+        ],
     },
     shard_count = {
         "cpu": 30,