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
This commit is contained in:
Peter Hawkins 2024-05-24 07:24:36 -07:00 committed by jax authors
parent be7939a20e
commit 1e01fa7b0f

View File

@ -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 = {