Revert "[dtensor] move all tests to distribute/tensor folder (#144166)"

This reverts commit 2e1ea85.

Reverted #144166 on behalf of https://github.com/huydhn due to Sorry for reverting your change, but inductor/test_compiled_autograd needs to be updated ([comment](#144166 (comment)))

Author: pytorchmergebot

.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

.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

.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/**

test/distributed/_tensor/README.md

@@ -0,0 +1,11 @@
+## Run distributed tensor tests:
+
+from root, run (either CPU or GPU)
+
+`pytest test/spmd/tensor/test_tensor.py`
+
+`pytest test/spmd/tensor/test_ddp.py`
+
+run specific test case and print stdout/stderr:
+
+`pytest test/spmd/tensor/test_tensor.py -s -k test_tensor_from_local`

test/distributed/tensor/__init__.py renamed to test/distributed/_tensor/__init__.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

test/distributed/tensor/debug/test_comm_mode.py renamed to test/distributed/_tensor/debug/test_comm_mode.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

test/distributed/tensor/debug/test_comm_mode_features.py renamed to test/distributed/_tensor/debug/test_comm_mode_features.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]