Enable test_xla_auto_sharding.py

pytorch · Mar 12, 2024 · 6ca8f97 · 6ca8f97
1 parent 55f1411
commit 6ca8f97
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 0 deletions.
diff --git a/test/run_tests.sh b/test/run_tests.sh
@@ -227,6 +227,7 @@ function run_xla_op_tests3 {
   run_test "$CDIR/spmd/test_xla_distributed_checkpoint.py"
   run_test "$CDIR/spmd/test_xla_spmd_python_api_interaction.py"
   run_test "$CDIR/spmd/test_dtensor_integration.py"
+  run_test "$CDIR/spmd/test_xla_auto_sharding.py"
   run_test "$CDIR/test_operations_hlo.py" "$@" --verbosity=$VERBOSITY
   run_test "$CDIR/test_input_output_aliases.py"
   run_test "$CDIR/test_torch_distributed_xla_backend.py"

diff --git a/test/spmd/test_dynamo_spmd.py b/test/spmd/test_dynamo_spmd.py
@@ -41,6 +41,7 @@ class DynamoSpmdInferenceTest(test_xla_sharding_base.XlaShardingTest):
 
   @classmethod
   def setUpClass(cls):
+    xr.use_spmd()
     super().setUpClass()
 
   def test_dynamo_spmd_basic(self):

diff --git a/test/spmd/test_spmd_graph_dump.py b/test/spmd/test_spmd_graph_dump.py
@@ -18,6 +18,7 @@ class SpmdGraphDumpTest(test_xla_sharding_base.XlaShardingTest):
 
   @classmethod
   def setUpClass(cls):
+    xr.use_spmd()
     super().setUpClass()
 
   def test_dump_with_output_sharding(self):

diff --git a/test/tpu/run_tests.sh b/test/tpu/run_tests.sh
@@ -10,6 +10,7 @@ python3 test/spmd/test_xla_virtual_device.py
 python3 test/spmd/test_xla_distributed_checkpoint.py
 python3 test/spmd/test_train_spmd_linear_model.py
 python3 test/spmd/test_xla_spmd_python_api_interaction.py
+python3 test/spmd/test_xla_auto_sharding.py
 XLA_EXPERIMENTAL=nonzero:masked_select python3 test/ds/test_dynamic_shape_models.py -v
 XLA_EXPERIMENTAL=nonzero:masked_select python3 test/ds/test_dynamic_shapes.py -v
 python3 test/test_autocast.py