Separate pallas test into a separate test shard

Author: ghpvnist

.github/workflows/_tpu_ci.yml

@@ -23,6 +23,7 @@ jobs:
           - test/tpu/run_tests.sh
           - test/tpu/run_expensive_test_1.sh
           - test/tpu/run_expensive_test_2.sh
+          - test/tpu/run_pallas_test.sh
     steps:
       - name: Checkout actions
         if: inputs.has_code_changes == 'true'

test/tpu/run_pallas_test.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+set -xue
+CDIR="$(cd "$(dirname "$0")" ; pwd -P)"
+TEST_CDIR="$(dirname "$CDIR")"
+
+# This test takes ~370 seconds to run
+python3 "$TEST_CDIR/test_pallas.py" -v
+
+# This test takes ~15 seconds to run
+python3 "$TEST_CDIR/test_pallas_spmd.py"
+
+# This test takes ~15 seconds to run
+XLA_DISABLE_FUNCTIONALIZATION=1 python3 "$TEST_CDIR/test_pallas_spmd.py"

test/tpu/run_tests.sh

@@ -81,9 +81,6 @@ run_test "$_TEST_DIR/test_assume_pure.py"
 run_test "$_TEST_DIR/test_assume_pure_spmd.py"
 run_test "$_TEST_DIR/test_as_stride_use_slice.py"
 run_xla_hlo_debug run_test "$_TEST_DIR/scan/test_scan_debug.py"
-run_test "$_TEST_DIR/test_pallas.py" -v
-run_test "$_TEST_DIR/test_pallas_spmd.py"
-XLA_DISABLE_FUNCTIONALIZATION=1 run_test "$_TEST_DIR/test_pallas_spmd.py"
 run_test "$_TEST_DIR/test_splash_attention.py"
 run_test "$_TEST_DIR/test_profiler_session.py"
 run_test "$_TEST_DIR/test_input_output_aliases.py"