From ef65ba8f32ee185ee8a473575fcad07e4cd28acc Mon Sep 17 00:00:00 2001
From: Adam Paszke <apaszke@google.com>
Date: Wed, 29 Nov 2023 05:47:13 -0800
Subject: [PATCH] Internal change

PiperOrigin-RevId: 586312803
---
 tests/BUILD | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/tests/BUILD b/tests/BUILD
index aedcc7ca0ac2..a496e4578fc7 100644
--- 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",
     ],