[reland] pass shape/stride during tensor unflatten (pytorch#117340)

Reland of pytorch#113547 as the previous
PR reverted bc of torch.compile symbolic shape issue. Since we now disabled tensor
unflatten with dynamo.disable, we should not hit this issue again

Pull Request resolved: pytorch#117340
Approved by: https://github.com/Skylion007
ghstack dependencies: pytorch#117336

Author: wanchaol

torch/distributed/_tensor/placement_types.py

@@ -457,6 +457,12 @@ def shape(self) -> torch.Size:
             raise ValueError("tensor_meta is not set")
         return self.tensor_meta.shape
 
+    @property
+    def stride(self) -> Tuple[int, ...]:
+        if self.tensor_meta is None:
+            raise ValueError("tensor_meta is not set")
+        return self.tensor_meta.stride
+
     @property
     def ndim(self) -> int:
         if self.tensor_meta is None:

torch/distributed/tensor/parallel/_data_parallel_utils.py

@@ -36,6 +36,8 @@ def _unflatten_tensor(tensor, spec, *, device_handle=None, compute_stream=None):
         spec.mesh,
         spec.placements,
         run_check=False,
+        shape=spec.shape,
+        stride=spec.stride,
     )
     if tensor.requires_grad:
         # only register the hook if the tensor requires grad