[SDY] fix JAX layouts tests for Shardy.

PiperOrigin-RevId: 697715276
jax-ml · Nov 18, 2024 · 0ed6eae · 0ed6eae
1 parent 70b05f6
commit 0ed6eae
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tests/BUILD b/tests/BUILD
@@ -267,6 +267,9 @@ jax_multiplatform_test(
     backend_tags = {
         "tpu": ["requires-mem:16g"],  # Under tsan on 2x2 this test exceeds the default 12G memory limit.
     },
+    enable_configs = [
+        "tpu_v3_2x2_shardy",
+    ],
     tags = ["multiaccelerator"],
     deps = [
         "//jax:experimental",