From 01b4cb6de05ddbd95cab0e1d151fdfeac5a8f364 Mon Sep 17 00:00:00 2001 From: Peter Hawkins Date: Thu, 30 May 2024 05:31:46 -0700 Subject: [PATCH] Bump memory limits for array_test and layout_test on TPU CI. These use more than our CI's default memory limit (12GB) when run under tsan. PiperOrigin-RevId: 638618718 --- tests/BUILD | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/BUILD b/tests/BUILD index 83cbc368f..ebc02e05f 100644 --- a/tests/BUILD +++ b/tests/BUILD @@ -269,6 +269,9 @@ jax_test( jax_test( name = "layout_test", srcs = ["layout_test.py"], + backend_tags = { + "tpu": ["requires-mem:16g"], # Under tsan on 2x2 this test exceeds the default 12G memory limit. + }, tags = ["multiaccelerator"], ) @@ -318,6 +321,9 @@ jax_test( jax_test( name = "array_test", srcs = ["array_test.py"], + backend_tags = { + "tpu": ["requires-mem:16g"], # Under tsan on 2x2 this test exceeds the default 12G memory limit. + }, tags = ["multiaccelerator"], deps = [ "//jax:experimental",