From efdc83d7241e15aaef925cd9e2c26b06bb703e58 Mon Sep 17 00:00:00 2001
From: Joshua Lang <joshualang@google.com>
Date: Mon, 24 Nov 2025 08:10:42 -0800
Subject: [PATCH] Update Jax B200 single gpu tests to use MIG partitions.

PiperOrigin-RevId: 836231103
---
 tests/pallas/BUILD | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tests/pallas/BUILD b/tests/pallas/BUILD
index 51be18640955..2baf5a24090c 100644
--- a/tests/pallas/BUILD
+++ b/tests/pallas/BUILD
@@ -250,11 +250,16 @@ jax_multiplatform_test(
     srcs = [
         "mosaic_gpu_test.py",
     ],
+    disable_configs = [
+        # TODO(b/462499936): Re-enable when test passes on MIG partition.
+        "gpu_b200",
+    ],
     enable_backends = [],
     enable_configs = [
         "gpu_h100_x32",
         "gpu_h100",
-        "gpu_b200",
+        # TODO(b/462499936): Remove gpu_b200_full when test passes onMIG partition.
+        "gpu_b200_full",
     ],
     shard_count = 4,
     tags = [