From 1e01fa7b0f1355c522f8420569dc778f2633c629 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Fri, 24 May 2024 07:24:36 -0700
Subject: [PATCH] Bump up memory requirements for pmap_test on TPU.

This test recently started exceeding the default memory requirement when run under tsan. I'm not entirely sure why, but perhaps some change pushed it just over our CI's 12GB default limit.

PiperOrigin-RevId: 636910434
---
 tests/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/BUILD b/tests/BUILD
index 54c77cd54..83cbc368f 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -710,6 +710,7 @@ jax_test(
     backend_tags = {
         "tpu": [
             "noasan",  # Times out under asan.
+            "requires-mem:16g",  # Under tsan on 2x2 this test exceeds the default 12G memory limit.
         ],
     },
     shard_count = {