From bea63892b98ed521fc793c9e33e125e256e47a83 Mon Sep 17 00:00:00 2001
From: Jiewen Tan <jwtan@google.com>
Date: Fri, 10 May 2024 01:21:04 +0000
Subject: [PATCH] add pallas_spmd.py to tpu ci

---
 test/tpu/run_tests.sh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/tpu/run_tests.sh b/test/tpu/run_tests.sh
index 2653d0ed8c2..34dbfe0db96 100755
--- a/test/tpu/run_tests.sh
+++ b/test/tpu/run_tests.sh
@@ -21,6 +21,7 @@ python3 test/pjrt/test_dtypes.py
 python3 test/pjrt/test_dynamic_plugin_tpu.py
 python3 test/test_fori_loop_with_while_loop_simple_add_dispatch_in_torch.py
 python3 test/test_pallas.py
+python3 test/test_pallas_spmd.py
 python3 test/test_input_output_aliases.py
 python3 test/test_megablox.py
 python3 test/torch_distributed/test_torch_distributed_all_gather_xla_backend.py