pytorch · wanchaol · Jan 3, 2025 · Jan 3, 2025 · Jan 3, 2025 · Jan 6, 2025
diff --git a/.ci/pytorch/multigpu-test.sh b/.ci/pytorch/multigpu-test.sh
@@ -43,8 +43,8 @@ if [[ "${SHARD_NUMBER:-2}" == "2" ]]; then
     time python test/run_test.py --verbose -i distributed/test_functional_api
 
     # DTensor tests
-    time python test/run_test.py --verbose -i distributed/_tensor/test_random_ops
-    time python test/run_test.py --verbose -i distributed/_tensor/test_dtensor_compile
+    time python test/run_test.py --verbose -i distributed/tensor/test_random_ops
+    time python test/run_test.py --verbose -i distributed/tensor/test_dtensor_compile
 
     # DeviceMesh test
     time python test/run_test.py --verbose -i distributed/test_device_mesh

diff --git a/.ci/pytorch/test.sh b/.ci/pytorch/test.sh
@@ -329,7 +329,7 @@ test_inductor_distributed() {
   python test/run_test.py -i inductor/test_aot_inductor.py -k test_non_default_cuda_device --verbose
   python test/run_test.py -i inductor/test_aot_inductor.py -k test_replicate_on_devices --verbose
   python test/run_test.py -i distributed/test_c10d_functional_native.py --verbose
-  python test/run_test.py -i distributed/_tensor/test_dtensor_compile.py --verbose
+  python test/run_test.py -i distributed/tensor/test_dtensor_compile.py --verbose
   python test/run_test.py -i distributed/tensor/parallel/test_micro_pipeline_tp.py --verbose
   python test/run_test.py -i distributed/_composable/test_replicate_with_compiler.py --verbose
   python test/run_test.py -i distributed/_composable/fsdp/test_fully_shard_comm.py --verbose

diff --git a/.github/labeler.yml b/.github/labeler.yml
@@ -30,9 +30,9 @@
 - torch/fx/experimental/sym_node.py
 - torch/fx/experimental/validator.py
 - torch/fx/experimental/proxy_tensor.py
-- test/distributed/_tensor/test_dtensor_compile.py
+- test/distributed/tensor/test_dtensor_compile.py
 - test/distributed/tensor/parallel/test_fsdp_2d_parallel.py
-- torch/distributed/_tensor/**
+- torch/distributed/tensor/**
 - torch/distributed/fsdp/**
 - torch/csrc/inductor/**
 - torch/csrc/dynamo/**

diff --git a/test/distributed/_tensor/README.md b/test/distributed/_tensor/README.md
diff --git a/test/distributed/tensor/README.md b/test/distributed/tensor/README.md
@@ -0,0 +1,10 @@
+## Run distributed tensor tests:
+
+from root, run (either CPU or GPU)
+
+`pytest test/distributed/tensor/test_dtensor.py`
+
+
+run specific test cases and print stdout/stderr:
+
+`pytest test/distributed/tensor/test_dtensor.py -s -k test_from_local`
diff --git a/test/distributed/_tensor/__init__.py → test/distributed/tensor/__init__.py b/test/distributed/_tensor/__init__.py → test/distributed/tensor/__init__.py
diff --git a/...stributed/_tensor/debug/test_comm_mode.py → ...istributed/tensor/debug/test_comm_mode.py b/...stributed/_tensor/debug/test_comm_mode.py → ...istributed/tensor/debug/test_comm_mode.py
diff --git a/.../_tensor/debug/test_comm_mode_features.py → ...d/tensor/debug/test_comm_mode_features.py b/.../_tensor/debug/test_comm_mode_features.py → ...d/tensor/debug/test_comm_mode_features.py
diff --git a/...ributed/_tensor/debug/test_op_coverage.py → ...tributed/tensor/debug/test_op_coverage.py b/...ributed/_tensor/debug/test_op_coverage.py → ...tributed/tensor/debug/test_op_coverage.py
diff --git a/...ed/_tensor/experimental/test_local_map.py → ...ted/tensor/experimental/test_local_map.py b/...ed/_tensor/experimental/test_local_map.py → ...ted/tensor/experimental/test_local_map.py
diff --git a/...or/experimental/test_register_sharding.py → ...or/experimental/test_register_sharding.py b/...or/experimental/test_register_sharding.py → ...or/experimental/test_register_sharding.py
diff --git a/..._tensor/experimental/test_tp_transform.py → .../tensor/experimental/test_tp_transform.py b/..._tensor/experimental/test_tp_transform.py → .../tensor/experimental/test_tp_transform.py
diff --git a/test/distributed/_tensor/test_api.py → test/distributed/tensor/test_api.py b/test/distributed/_tensor/test_api.py → test/distributed/tensor/test_api.py
diff --git a/test/distributed/_tensor/test_attention.py → test/distributed/tensor/test_attention.py b/test/distributed/_tensor/test_attention.py → test/distributed/tensor/test_attention.py
diff --git a/.../distributed/_tensor/test_common_rules.py → test/distributed/tensor/test_common_rules.py b/.../distributed/_tensor/test_common_rules.py → test/distributed/tensor/test_common_rules.py
diff --git a/...stributed/_tensor/test_convolution_ops.py → ...istributed/tensor/test_convolution_ops.py b/...stributed/_tensor/test_convolution_ops.py → ...istributed/tensor/test_convolution_ops.py
diff --git a/test/distributed/_tensor/test_dtensor.py → test/distributed/tensor/test_dtensor.py b/test/distributed/_tensor/test_dtensor.py → test/distributed/tensor/test_dtensor.py
diff --git a/...stributed/_tensor/test_dtensor_compile.py → ...istributed/tensor/test_dtensor_compile.py b/...stributed/_tensor/test_dtensor_compile.py → ...istributed/tensor/test_dtensor_compile.py
diff --git a/test/distributed/_tensor/test_dtensor_ops.py → test/distributed/tensor/test_dtensor_ops.py b/test/distributed/_tensor/test_dtensor_ops.py → test/distributed/tensor/test_dtensor_ops.py
@@ -82,7 +82,7 @@ def wrapped(fn):
 
 # Re-generate this failed list, turn on dry_run of the below func
 # check_dtensor_func(self, test, op, dry_run=True), then run sth
-# like python test/distributed/_tensor/test_dtensor_ops.py > failed.expect
+# like python test/distributed/tensor/test_dtensor_ops.py > failed.expect
 dtensor_fails = {
     # these sometimes pass and sometimes fail
     # we need to remove many of them from list once op

diff --git a/...distributed/_tensor/test_embedding_ops.py → .../distributed/tensor/test_embedding_ops.py b/...distributed/_tensor/test_embedding_ops.py → .../distributed/tensor/test_embedding_ops.py
diff --git a/...tributed/_tensor/test_experimental_ops.py → ...stributed/tensor/test_experimental_ops.py b/...tributed/_tensor/test_experimental_ops.py → ...stributed/tensor/test_experimental_ops.py
diff --git a/test/distributed/_tensor/test_init.py → test/distributed/tensor/test_init.py b/test/distributed/_tensor/test_init.py → test/distributed/tensor/test_init.py
diff --git a/test/distributed/_tensor/test_math_ops.py → test/distributed/tensor/test_math_ops.py b/test/distributed/_tensor/test_math_ops.py → test/distributed/tensor/test_math_ops.py
diff --git a/test/distributed/_tensor/test_matrix_ops.py → test/distributed/tensor/test_matrix_ops.py b/test/distributed/_tensor/test_matrix_ops.py → test/distributed/tensor/test_matrix_ops.py
diff --git a/test/distributed/_tensor/test_op_strategy.py → test/distributed/tensor/test_op_strategy.py b/test/distributed/_tensor/test_op_strategy.py → test/distributed/tensor/test_op_strategy.py
diff --git a/test/distributed/_tensor/test_optimizers.py → test/distributed/tensor/test_optimizers.py b/test/distributed/_tensor/test_optimizers.py → test/distributed/tensor/test_optimizers.py
diff --git a/...distributed/_tensor/test_pointwise_ops.py → .../distributed/tensor/test_pointwise_ops.py b/...distributed/_tensor/test_pointwise_ops.py → .../distributed/tensor/test_pointwise_ops.py
diff --git a/test/distributed/_tensor/test_random_ops.py → test/distributed/tensor/test_random_ops.py b/test/distributed/_tensor/test_random_ops.py → test/distributed/tensor/test_random_ops.py
diff --git a/.../distributed/_tensor/test_redistribute.py → test/distributed/tensor/test_redistribute.py b/.../distributed/_tensor/test_redistribute.py → test/distributed/tensor/test_redistribute.py
diff --git a/test/distributed/_tensor/test_tensor_ops.py → test/distributed/tensor/test_tensor_ops.py b/test/distributed/_tensor/test_tensor_ops.py → test/distributed/tensor/test_tensor_ops.py
diff --git a/test/distributed/_tensor/test_utils.py → test/distributed/tensor/test_utils.py b/test/distributed/_tensor/test_utils.py → test/distributed/tensor/test_utils.py
diff --git a/test/distributed/_tensor/test_view_ops.py → test/distributed/tensor/test_view_ops.py b/test/distributed/_tensor/test_view_ops.py → test/distributed/tensor/test_view_ops.py
diff --git a/...stributed/_tensor/test_xla_integration.py → ...istributed/tensor/test_xla_integration.py b/...stributed/_tensor/test_xla_integration.py → ...istributed/tensor/test_xla_integration.py
diff --git a/test/inductor/test_compiled_autograd.py b/test/inductor/test_compiled_autograd.py
@@ -3629,7 +3629,7 @@ def wrap_test_class(orig_cls):
 TestAutogradWithCompiledAutograd = wrap_test_class(test_autograd.TestAutograd)
 TestCustomOpWithCompiledAutograd = wrap_test_class(test_custom_ops.TestCustomOp)
 if torch.distributed.is_available() and HAS_CUDA:
-    test_dtensor = load_test_module("distributed/_tensor/test_dtensor_compile")
+    test_dtensor = load_test_module("distributed/tensor/test_dtensor_compile")
     TestDTensorCompileWithCompiledAutograd = wrap_test_class(
         test_dtensor.TestDTensorCompile
     )

diff --git a/test/run_test.py b/test/run_test.py
@@ -182,7 +182,7 @@ def __contains__(self, item):
     "test_jit_legacy",
     "test_cuda_nvml_based_avail",
     "test_jit_cuda_fuser",
-    "distributed/_tensor/test_attention",
+    "distributed/tensor/test_attention",
 ]
 
 # whitelist of tests for s390x

diff --git a/torch/fx/experimental/symbolic_shapes.py b/torch/fx/experimental/symbolic_shapes.py
@@ -837,7 +837,7 @@ def free_symbols(val: IterateExprs) -> OrderedSet[sympy.Symbol]:
         return OrderedSet()
 
     # TODO: Apparently, returning an OrderedSet here breaks
-    # python test/distributed/_tensor/test_dtensor_compile.py TestDTensorCompile.test_dtensor_dynamic
+    # python test/distributed/tensor/test_dtensor_compile.py TestDTensorCompile.test_dtensor_dynamic
     return first_expr.free_symbols.union(*(e.free_symbols for e in itr))  # type: ignore[return-value]