Internal change

PiperOrigin-RevId: 586312803
2025-04-15 19:36:06 +00:00 · 2023-11-29 05:47:13 -08:00 · 2023-11-29 05:47:13 -08:00 · ef65ba8f32
commit ef65ba8f32
parent 458a8962be
1 changed files with 8 additions and 3 deletions
--- a/tests/BUILD
+++ b/tests/BUILD
@ -172,7 +172,7 @@ jax_test(
    shard_count = {
        "cpu": 10,
        "gpu": 10,
-        "tpu": 20,
+        "tpu": 40,
    },
 )

@ -403,7 +403,7 @@ jax_test(
    shard_count = {
        "cpu": 30,
        "gpu": 30,
-        "tpu": 20,
+        "tpu": 40,
    },
 )

@ -1223,7 +1223,12 @@ jax_test(
        "gpu": 10,
        "tpu": 50,
    },
-    tags = ["multiaccelerator"],
+    tags = [
+        "multiaccelerator",
+        "noasan",
+        "nomsan",
+        "notsan",
+    ],  # Times out under *SAN.
    deps = [
        "//jax:tree_util",
    ],