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"