From 0ed6eaeb4a0c5ebf7679f3877b01bd7d6df29bae Mon Sep 17 00:00:00 2001
From: Bill Varcho <varcho@google.com>
Date: Mon, 18 Nov 2024 12:13:55 -0800
Subject: [PATCH] [SDY] fix JAX layouts tests for Shardy.

PiperOrigin-RevId: 697715276
---
 tests/BUILD | 3 +++
 1 file changed, 3 insertions(+)

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