patched_pytorch_v2.2.1_rc3.patch

diff --git a/aten/src/ATen/FunctionalInverses.cpp b/aten/src/ATen/FunctionalInverses.cpp
index af0e5af3be8..9896f16a84e 100644
--- a/aten/src/ATen/FunctionalInverses.cpp
+++ b/aten/src/ATen/FunctionalInverses.cpp
@@ -151,6 +151,12 @@ Tensor FunctionalInverses::expand_copy_inverse(const Tensor& base, const Tensor&
     return at::sum_to(mutated_view, base.sym_sizes(),/*always_return_non_view=*/!reapply_views);
 }

+Tensor FunctionalInverses::expand_as_copy_inverse(const Tensor& base, const Tensor& mutated_view, bool reapply_views,const Tensor& other) {
+    return at::sum_to(mutated_view, base.sym_sizes(),/*always_return_non_view=*/!reapply_views);
+}
+
+
+
 Tensor FunctionalInverses::permute_copy_inverse(const Tensor& base, const Tensor& mutated_view, bool reapply_views, at::IntArrayRef dims) {
     return at::functionalization::permute_copy_inverse(mutated_view, dims, reapply_views);
 }
diff --git a/aten/src/ATen/functorch/BatchRulesDecompositions.cpp b/aten/src/ATen/functorch/BatchRulesDecompositions.cpp
index 1b179a505e9..b1beaa67ae7 100644
--- a/aten/src/ATen/functorch/BatchRulesDecompositions.cpp
+++ b/aten/src/ATen/functorch/BatchRulesDecompositions.cpp
@@ -296,7 +296,7 @@ TORCH_LIBRARY_IMPL(aten, FuncTorchBatchedDecomposition, m) {
   OP_DECOMPOSE2(trapz, x);
   OP_DECOMPOSE2(trapz, dx);
   OP_DECOMPOSE(unsafe_chunk);
-  m.impl("value_selecting_reduction_backward", native::value_selecting_reduction_backward_symint);
+  m.impl("value_selecting_reduction_backward", native::value_selecting_reduction_backward);
   OP_DECOMPOSE(var);
   OP_DECOMPOSE2(var, dim);
   OP_DECOMPOSE(var_mean);
diff --git a/aten/src/ATen/native/Onehot.cpp b/aten/src/ATen/native/Onehot.cpp
index 41b7a696186..26fd0979c39 100644
--- a/aten/src/ATen/native/Onehot.cpp
+++ b/aten/src/ATen/native/Onehot.cpp
@@ -5,7 +5,9 @@
 #include <ATen/Functions.h>
 #include <ATen/NativeFunctions.h>
 #else
+#include <ATen/ops/arange.h>
 #include <ATen/ops/empty.h>
+#include <ATen/ops/eq.h>
 #include <ATen/ops/one_hot_native.h>
 #include <ATen/ops/zeros.h>
 #endif
@@ -14,6 +16,17 @@ namespace at { namespace native {

 Tensor one_hot(const Tensor &self, int64_t num_classes) {
     TORCH_CHECK(self.dtype() == kLong, "one_hot is only applicable to index tensor.");
+    // using meta bit test to catch Fake Tensor as well until __torch_function__
+    if (self.key_set().has_all(DispatchKeySet(BackendComponent::MetaBit)) ||
+            self.key_set().has_all(DispatchKeySet(DispatchKey::Python))) {
+        // functional version that torch.compiles better and works with dynamic shapes
+        if (num_classes == -1) {
+          num_classes = self.max().item().toLong() + 1;
+        }
+        at::Tensor index = at::arange(num_classes, self.options());
+        return at::eq(self.unsqueeze(-1), index).to(kLong);
+    }
+
     auto shape = self.sizes().vec();

     // empty tensor could be converted to one hot representation,
diff --git a/aten/src/ATen/native/ReduceOps.cpp b/aten/src/ATen/native/ReduceOps.cpp
index 7a47490c674..a2c54db9424 100644
--- a/aten/src/ATen/native/ReduceOps.cpp
+++ b/aten/src/ATen/native/ReduceOps.cpp
@@ -2228,26 +2228,21 @@ bool cpu_equal(const Tensor& self, const Tensor& other) {
   return result.load();
 }

-static Tensor value_selecting_reduction_backward(const Tensor& grad, int64_t dim, const Tensor& indices, at::IntArrayRef sizes, bool keepdim) {
-    return at::native::value_selecting_reduction_backward_symint(grad, dim, indices, c10::fromIntArrayRefSlow(sizes), keepdim);
-}
-
-
 // max(dim), min(dim), topk(dim), mode(dim), are examples of reduction
 // functions that select values. value_selecting_reduction_backward is the
 // backward function for those operators; it propagates the grad to the
 // specific value locations referred to at `indices`.
-Tensor value_selecting_reduction_backward_symint(const Tensor& grad, int64_t dim, const Tensor& indices, c10::SymIntArrayRef sizes, bool keepdim) {
+Tensor value_selecting_reduction_backward(const Tensor& grad, int64_t dim, const Tensor& indices, const Tensor& src, bool keepdim) {
   auto inplace_scatter_if_not_tensor_subclass =
       [&](const Tensor& grad_out, const Tensor& indices_) {
-        auto grad_in = at::zeros_symint(sizes, grad_out.options());
+        auto grad_in = at::zeros_like(src, grad_out.options());
         if (areAnyTensorSubclassLike({grad, indices})) {
           return grad_in.scatter(dim, indices_, grad_out);
         }
         return grad_in.scatter_(dim, indices_, grad_out);
       };

-  if (!keepdim && !sizes.empty()) {
+  if (!keepdim && !src.sizes().empty()) {
     auto grad_ = grad.unsqueeze(dim);
     auto indices_ = indices.unsqueeze(dim);
     return inplace_scatter_if_not_tensor_subclass(grad_, indices_);
diff --git a/aten/src/ATen/native/TensorShape.cpp b/aten/src/ATen/native/TensorShape.cpp
index 0a018fbc8db..a5e4643ae53 100644
--- a/aten/src/ATen/native/TensorShape.cpp
+++ b/aten/src/ATen/native/TensorShape.cpp
@@ -109,6 +109,7 @@
 #include <ATen/ops/empty_quantized.h>
 #include <ATen/ops/expand_as_native.h>
 #include <ATen/ops/expand_copy_native.h>
+#include <ATen/ops/expand_as_copy_native.h>
 #include <ATen/ops/expand_native.h>
 #include <ATen/ops/flatten_dense_tensors_native.h>
 #include <ATen/ops/flatten_native.h>
@@ -1143,7 +1144,21 @@ Tensor expand(const Tensor& self, c10::IntArrayRef size, bool /*unused*/) {
 }

 Tensor expand_as(const Tensor& self, const Tensor& other) {
-  return self.expand_symint(other.sym_sizes());
+  IntArrayRef size = other.sizes();
+  TORCH_CHECK(size.size() >= (size_t)self.dim(),
+           "expand(", self.toString(), "{", self.sizes(), "}, size=", size,
+           "): the number of sizes provided (", size.size(), ") ",
+           "must be greater or equal to the number of dimensions in the tensor (",
+           self.dim(), ")");
+  TORCH_CHECK(!self.is_sparse() && !at::sparse_csr::is_sparse_compressed(self),
+            "expand is unsupported for ", self.layout(), " tensors");
+
+  auto expandedSizesAndStrides = inferExpandGeometry_dimvector(self.sizes(), self.strides(), size);
+
+  auto result = self.as_strided(
+      expandedSizesAndStrides.sizes, expandedSizesAndStrides.strides);
+  namedinference::propagate_names_for_expand(result, self);
+  return result;
 }

 Tensor sum_to_size_symint(const Tensor& self, SymIntArrayRef size) {
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
index 35a1049e209..604f53ac734 100644
--- a/aten/src/ATen/native/native_functions.yaml
+++ b/aten/src/ATen/native/native_functions.yaml
@@ -2595,6 +2595,8 @@
   variants: method  # This is method-only to match the previous tensor API. In the future we could make this a function too.
   device_check: NoCheck
   device_guard: False
+  dispatch:
+    CompositeExplicitAutograd: expand_as

 # decomposes to eye.m
 - func: eye(SymInt n, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
@@ -3759,12 +3761,10 @@
 - func: max.names_dim_max(Tensor self, Dimname dim, bool keepdim=False, *, Tensor(a!) max, Tensor(b!) max_values) -> (Tensor(a!) values, Tensor(b!) indices)
   device_check: NoCheck   # TensorIterator

-- func: value_selecting_reduction_backward(Tensor grad, int dim, Tensor indices, SymInt[] sizes, bool keepdim) -> Tensor
+- func: value_selecting_reduction_backward(Tensor grad, int dim, Tensor indices, Tensor sizes, bool keepdim) -> Tensor
   variants: function
   device_check: NoCheck
   device_guard: False
-  dispatch:
-    CompositeImplicitAutograd: value_selecting_reduction_backward_symint

 - func: amax(Tensor self, int[1] dim=[], bool keepdim=False) -> Tensor
   variants: function, method
@@ -14225,6 +14225,13 @@
   tags: view_copy
   autogen: expand_copy.out

+- func: expand_as_copy(Tensor self, Tensor other) -> Tensor
+  variants: function
+  dispatch:
+    CompositeExplicitAutogradNonFunctional: expand_as_copy
+  tags: view_copy
+  autogen: expand_as_copy.out
+
 - func: permute_copy(Tensor self, int[] dims) -> Tensor
   variants: function
   dispatch:
diff --git a/test/distributed/_tensor/test_dtensor.py b/test/distributed/_tensor/test_dtensor.py
index a83efe539e4..e190c5b97d5 100644
--- a/test/distributed/_tensor/test_dtensor.py
+++ b/test/distributed/_tensor/test_dtensor.py
@@ -109,6 +109,16 @@ class DTensorTest(DTensorTestBase):
             value_tensor = torch.empty_like(meta_dtensor.to_local()).fill_(1.5)
             self.assertEqual(meta_dtensor.to_local(), value_tensor)

+    @with_comms
+    def test_dtensor_local_tensor_storage(self):
+        device_mesh = self.build_device_mesh()
+        shard0_spec = [Shard(0)]
+        local_tensor = torch.randn(4, 8)
+        dist_tensor = DTensor.from_local(local_tensor, device_mesh, shard0_spec)
+        self.assertEqual(dist_tensor.data_ptr(), dist_tensor._local_tensor.data_ptr())
+        local_tensor = dist_tensor.to_local()
+        self.assertEqual(dist_tensor.data_ptr(), local_tensor.data_ptr())
+
     @with_comms
     def test_modules_w_meta_dtensor(self):
         model = DummyMLP("meta")
diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml
index 2c6886a36cc..6d651249354 100644
--- a/tools/autograd/derivatives.yaml
+++ b/tools/autograd/derivatives.yaml
@@ -892,7 +892,7 @@
   self: non_differentiable

 - name: kthvalue(Tensor self, int k, int dim=-1, bool keepdim=False) -> (Tensor values, Tensor indices)
-  self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), keepdim)
+  self: value_selecting_reduction_backward(grad, dim, indices, self, keepdim)
   values: gather_with_keepdimed_indices(self_t, dim, indices, keepdim)

 - name: le_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)
@@ -1084,7 +1084,7 @@
   result: linalg_matrix_exp_differential(self_p, self_t, /*adjoint*/ false)

 - name: max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)
-  self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), keepdim)
+  self: value_selecting_reduction_backward(grad, dim, indices, self, keepdim)
   values: gather_with_keepdimed_indices(self_t, dim, indices, keepdim)

 - name: max(Tensor self) -> Tensor
@@ -1132,15 +1132,15 @@
 # The backward implementation is correct in the sense that it returns the
 # subgradient on one side.
 - name: median.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)
-  self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), keepdim)
+  self: value_selecting_reduction_backward(grad, dim, indices, self, keepdim)
   values: gather_with_keepdimed_indices(self_t, dim, indices, keepdim)

 - name: nanmedian.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)
-  self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), keepdim)
+  self: value_selecting_reduction_backward(grad, dim, indices, self, keepdim)
   values: gather_with_keepdimed_indices(self_t, dim, indices, keepdim)

 - name: min.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)
-  self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), keepdim)
+  self: value_selecting_reduction_backward(grad, dim, indices, self, keepdim)
   values: gather_with_keepdimed_indices(self_t, dim, indices, keepdim)

 - name: min(Tensor self) -> Tensor
@@ -1171,7 +1171,7 @@
   result: at::mm(self_t, mat2_p) + at::mm(self_p, mat2_t)

 - name: mode(Tensor self, int dim=-1, bool keepdim=False) -> (Tensor values, Tensor indices)
-  self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), keepdim)
+  self: value_selecting_reduction_backward(grad, dim, indices, self, keepdim)
   values: gather_with_keepdimed_indices(self_t, dim, indices, keepdim)

 - name: mul.Tensor(Tensor self, Tensor other) -> Tensor
@@ -1526,12 +1526,12 @@
   output_differentiability: [True, False, False, False]  # LU is an auxiliary tensor not exposed to the user

 - name: sort(Tensor self, int dim=-1, bool descending=False) -> (Tensor values, Tensor indices)
-  self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), true)
+  self: value_selecting_reduction_backward(grad, dim, indices, self, true)
   output_differentiability: [True, False]
   values: gather_with_keepdimed_indices(self_t, dim, indices, true)

 - name: sort.stable(Tensor self, *, bool? stable, int dim=-1, bool descending=False) -> (Tensor values, Tensor indices)
-  self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), true)
+  self: value_selecting_reduction_backward(grad, dim, indices, self, true)
   output_differentiability: [True, False]
   values: gather_with_keepdimed_indices(self_t, dim, indices, true)

@@ -1692,7 +1692,7 @@
   result: auto_element_wise

 - name: topk(Tensor self, SymInt k, int dim=-1, bool largest=True, bool sorted=True) -> (Tensor values, Tensor indices)
-  self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), true)
+  self: value_selecting_reduction_backward(grad, dim, indices, self, true)
   output_differentiability: [True, False]
   values: gather(self_t, dim, indices)

diff --git a/tools/autograd/gen_inplace_or_view_type.py b/tools/autograd/gen_inplace_or_view_type.py
index ee1075cbed9..fea1c399012 100644
--- a/tools/autograd/gen_inplace_or_view_type.py
+++ b/tools/autograd/gen_inplace_or_view_type.py
@@ -315,6 +315,7 @@ def get_view_info(f: NativeFunction) -> Optional[str]:
 def emit_view_call(
     f: NativeFunction, input_base: str, unpacked_args: Sequence[str]
 ) -> str:
+
     # View replay functions use the standard Dispatcher::call API.
     return CALL_DISPATCH.substitute(
         unambiguous_name=f.func.name.unambiguous_name(), unpacked_args=unpacked_args
@@ -368,6 +369,7 @@ def emit_view_lambda(f: NativeFunction, unpacked_bindings: List[Binding]) -> str
             updated_unpacked_args.append(arg_value)
         elif (
             arg == "nested_size_" or arg == "nested_strides_" or arg == "offsets_"
+            or arg == "other_"
         ) and arg_type == ConstRefCType(BaseCType(tensorT)):
             # [NOTE] [Nested Arg Types]
             # This is temporary. Nested tensors will be migrating to use SymInts and
diff --git a/torch/_dynamo/variables/distributed.py b/torch/_dynamo/variables/distributed.py
index 54ad1cdf9b4..47605d96aed 100644
--- a/torch/_dynamo/variables/distributed.py
+++ b/torch/_dynamo/variables/distributed.py
@@ -24,9 +24,7 @@ class DistributedVariable(VariableTracker):
 def is_from_local(value):
     if not DistributedVariable.is_available():
         return False
-    from torch.distributed._tensor import DTensor
-
-    return inspect.isfunction(value) and value is DTensor.from_local
+    return inspect.isfunction(value) and value.__name__ == "from_local"


 def is_constant_pg_functions(value):
@@ -57,17 +55,17 @@ class PlacementClassVariable(DistributedVariable):
         if not DistributedVariable.is_available():
             return False

-        from torch.distributed._tensor.placement_types import Placement
+            if not isinstance(value, type):
+                return False
+            return value.__name__ in ("Placement", "Replicate", "Shard", "_Partial" "Partial", "InterleavedShard")

-        return type(value) is type and issubclass(value, Placement)
+    def as_python_constant(self):
+        return self.value

     def call_function(
         self, tx, args: "List[VariableTracker]", kwargs: "Dict[str, VariableTracker]"
     ) -> "VariableTracker":
-        if (
-            inspect.getattr_static(self.value, "__new__", None) in (object.__new__,)
-            and self.source
-        ):
+        if inspect.getattr_static(self.value, "__new__", None) in (object.__new__,) and self.source:
             # NOTE: we don't need to track mutations to the placement class as they
             # suppose to be immutable.
             new_obj = object.__new__(self.value)
@@ -90,9 +88,7 @@ class PlacementVariable(DistributedVariable):
         if not DistributedVariable.is_available():
             return False

-        from torch.distributed._tensor.placement_types import Placement
-
-        return isinstance(value, Placement)
+        return type(value).__name__ in ("Placement", "Replicate", "Shard", "_Partial" "Partial", "InterleavedShard")

     def as_python_constant(self):
         return self.value
@@ -106,15 +102,30 @@ class PlacementVariable(DistributedVariable):
     ) -> "VariableTracker":
         from . import ConstantVariable

-        allowed_methods = ["__init__", "__setattr__"]
-        # placement types dynamo tracking allows only __init__
-        # and __setattr__ methods, the latter is for case like `Shard(dim)`
-        if name in allowed_methods:
+# Placement types dynamo tracking only allows following methods
+        # and __setattr__  is for case like `shard(dim)` and methods.
+        # methods in the list must satisfy:
+        #    1. input arguments are constants and do not need to be guarded on;
+        #    2. output is constant with respect to their inputs
+        constant_fold_functions = [
+            "__init__",
+            "__setattr__",
+            "is_shard",
+            "is_partial",
+            "is_replicate",
+            "is_interleaved_shard",
+        ]
+        return_constant_functions = [
+            "is_shard",
+            "is_partial",
+            "is_replicate",
+            "is_interleaved_shard",
+        ]
+
+        if name in constant_fold_functions:
             try:
                 value_type = type(self.value)
-                assert (
-                    inspect.getattr_static(value_type, "__getattr__", None) is None
-                ), "no custom getattr allowed!"
+                assert inspect.getattr_static(value_type, "__getattr__", None) is None, "no custom getattr allowed!"
                 method = inspect.getattr_static(value_type, name)
             except AttributeError:
                 method = None
@@ -123,7 +134,9 @@ class PlacementVariable(DistributedVariable):

             args = [x.as_python_constant() for x in args]
             kwargs = {k: v.as_python_constant() for k, v in kwargs.items()}
-            method(self.value, *args, **kwargs)
+            out = method(self.value, *args, **kwargs)
+            if name in return_constant_functions:
+                return ConstantVariable(out)
             return self

         return super().call_method(tx, name, args, kwargs)
@@ -140,9 +153,7 @@ class DeviceMeshVariable(DistributedVariable):
         if not DistributedVariable.is_available():
             return False

-        from torch.distributed.device_mesh import DeviceMesh
-
-        return istype(value, DeviceMesh)
+        return type(value).__name__ == "DeviceMesh"

     def as_python_constant(self):
         return self.value
@@ -150,6 +161,9 @@ class DeviceMeshVariable(DistributedVariable):
     def var_getattr(self, tx, name: str) -> VariableTracker:
         if name == "ndim":
             return ConstantVariable.create(self.value.ndim)
+        if name == "device_type":
+            return ConstantVariable.create(self.value.device_type)
+
         return super().var_getattr(tx, name)


@@ -198,9 +212,7 @@ class ProcessGroupVariable(DistributedVariable):

     def var_getattr(self, tx, name):
         if name in ["rank", "size"]:
-            return variables.LambdaVariable(
-                lambda *args, **kwargs: self.call_method(tx, name, args, kwargs)
-            )
+            return variables.LambdaVariable(lambda *args, **kwargs: self.call_method(tx, name, args, kwargs))
         # TODO should this just raise unimplemented?
         return super().var_getattr(tx, name)

diff --git a/torch/_dynamo/variables/misc.py b/torch/_dynamo/variables/misc.py
index e5cf6f66730..755e28f331b 100644
--- a/torch/_dynamo/variables/misc.py
+++ b/torch/_dynamo/variables/misc.py
@@ -266,6 +266,64 @@ class NewGlobalVariable(VariableTracker):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)

+class BoundArgumentsVariable(VariableTracker):
+    """
+    This class is used to hack python code about `inspect` package, and not well-designed.
+    Please use it with caution.
+    """
+    def __init__(self, value, **kwargs):
+        super().__init__(**kwargs)
+        self.value = value
+        self.value.arguments = BoundArgumentsVariable.convert_to_variable_tracker(value.arguments)
+
+    @staticmethod
+    def convert_to_variable_tracker(arguments):
+        for key in arguments.keys():
+            val = arguments[key]
+            if isinstance(val, VariableTracker):
+                continue
+            # we only convert values of the first depth to VariableTracker,
+            # VariableTracker values of more than one depth should be handled by the caller.
+
+            if isinstance(val, (bool, int, float, type(None))):
+                arguments[key] = variables.ConstantVariable.create(val)
+            elif isinstance(val, set):
+                arguments[key] = variables.SetVariable(list(val))
+            elif isinstance(val, list):
+                arguments[key] = variables.ListVariable(val)
+            elif isinstance(val, dict):
+                arguments[key] = variables.ConstDictVariable(val)
+            elif isinstance(val, tuple):
+                arguments[key] = variables.TupleVariable(list(val))
+            else:
+                raise TypeError("unsupported arguments value type")
+        return arguments
+
+    @staticmethod
+    def create(value, **kwargs):
+        if kwargs:
+            unimplemented(f"inspect.BoundArguments with {kwargs}")
+        return BoundArgumentsVariable(value)
+
+    def var_getattr(self, tx, name: str) -> "VariableTracker":
+        if name == "arguments":
+            return variables.ConstDictVariable(self.value.arguments)
+        if name in ["signature", "_signature"]:
+            return InspectSignatureVariable.create(self.value._signature)
+        if name in ["args"]:
+            return variables.TupleVariable(list(self.value.args))
+        if name in ["kwargs"]:
+            return variables.ConstDictVariable(self.value.kwargs)
+        return super().var_getattr(tx, name)
+
+    def call_method(self, tx, name, args: List[VariableTracker], kwargs: Dict[str, VariableTracker]) -> VariableTracker:
+        if name == "apply_defaults":
+            assert len(args) == 0 and len(kwargs) == 0
+            self.value.apply_defaults()
+            BoundArgumentsVariable.convert_to_variable_tracker(self.value.arguments)
+            return variables.ConstantVariable.create(None)
+        return super().call_method(tx, name, args, kwargs)
+

 class InspectSignatureVariable(VariableTracker):
     """represents inspect.signature(...)"""
@@ -279,23 +337,52 @@ class InspectSignatureVariable(VariableTracker):
     def __init__(self, inspected: VariableTracker, **kwargs):
         super().__init__(**kwargs)
         self.inspected = inspected
+        self.python_signature = None
+        if isinstance(self.inspected, variables.UserMethodVariable):
+            self.python_signature = inspect.signature(getattr(self.inspected.obj.value, self.inspected.fn.__name__))
+        elif isinstance(self.inspected, variables.UserFunctionVariable):
+            self.python_signature = inspect.signature(self.inspected.fn)
+        else:
+            unimplemented("unsupported callable")

     def var_getattr(self, tx, name: str) -> "VariableTracker":
         if name == "parameters":
+            paramters = self.python_signature.parameters
             return variables.ConstDictVariable(
                 {
-                    name: InspectParameterVariable()
-                    for name in self.inspected.inspect_parameter_names()
+                    variables.ConstantVariable.create(name): InspectParameterVariable(value)
+                    for name, value in paramters.items()
                 },
                 user_cls=dict,
             )
         return super().var_getattr(tx, name)

+    def call_method(self, tx, name, args: List[VariableTracker], kwargs: Dict[str, VariableTracker]) -> VariableTracker:
+        if name == "bind":
+            # NOTE: InspectSignatureVariable only record the inspected user_method or function
+            # we need recover from it.
+            return BoundArgumentsVariable.create(self.python_signature.bind(*args, **kwargs))
+        return super().call_method(tx, name, args, kwargs)
+

 class InspectParameterVariable(VariableTracker):
     """This is not implemented, if used will graph break."""
+    def __init__(self, value, **kwargs):
+        super().__init__(**kwargs)
+        self.value = value

-    pass
+    @staticmethod
+    def create(value, **kwargs):
+        if kwargs:
+            unimplemented(f"inspect.signature with {kwargs}")
+        return InspectParameterVariable(value=value)
+
+    def var_getattr(self, tx, name: str) -> "VariableTracker":
+        if name in ["POSITIONAL_ONLY", "POSITIONAL_OR_KEYWORD", "VAR_POSITIONAL", "KEYWORD_ONLY", "VAR_KEYWORD"]:
+            return variables.ConstantVariable.create(getattr(inspect._ParameterKind, name))
+        if name in ["kind", "name", "default"]:
+            return variables.ConstantVariable.create(getattr(self.value, name))
+        return super().var_getattr(tx, name)


 def produce_trampoline_autograd_fwd(fn_cls):
diff --git a/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py b/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py
index 16eef07af02..ce82a2675d4 100644
--- a/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py
+++ b/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py
@@ -102,9 +102,10 @@ def aot_dispatch_base_graph(

     # TODO: should factor this into a separate function for export that always only returns just the graph.
     if aot_config.is_export:
-        assert (
-            maybe_subclass_meta is None
-        ), "aot_export_module does not support tensor subclass inputs for now."
+        # NOTE: hack, do not check if traced tenosr is a subclass
+        # assert (
+        #     maybe_subclass_meta is None
+        # ), "aot_export_module does not support tensor subclass inputs for now."
         return fw_module
     return fw_module, list(updated_flat_args_subclasses_desugared), maybe_subclass_meta

diff --git a/torch/_functorch/_aot_autograd/subclass_utils.py b/torch/_functorch/_aot_autograd/subclass_utils.py
index 0514c1c4d56..4d813fe64b5 100644
--- a/torch/_functorch/_aot_autograd/subclass_utils.py
+++ b/torch/_functorch/_aot_autograd/subclass_utils.py
@@ -16,6 +16,27 @@ from .utils import strict_zip

 zip = strict_zip

+def is_dtensor_subclass_dispatch(args, fw_metadata: ViewAndMutationMeta) -> bool:
+    args_flattened = pytree.arg_tree_leaves(*args)
+    # NOTE: hack: separately check DTensor dispatch
+    any_subclass_args = all(
+        (
+            is_traceable_wrapper_subclass(x) and
+            x.__class__.__name__ == "DTensor"
+        )
+        for x in args_flattened
+        if isinstance(x, Tensor)
+    )
+    any_subclass_outputs = all(
+        (
+            is_traceable_wrapper_subclass(x) and
+            x.__class__.__name__ == "DTensor"
+        )
+        for x in fw_metadata.traced_tangents
+        if isinstance(x, Tensor)
+    )
+    return any_subclass_args or any_subclass_outputs
+

 def requires_subclass_dispatch(args, fw_metadata: ViewAndMutationMeta) -> bool:
     args_flattened = pytree.arg_tree_leaves(*args)
diff --git a/torch/_functorch/aot_autograd.py b/torch/_functorch/aot_autograd.py
index 837fe2ab4b6..b38b2c2bedc 100644
--- a/torch/_functorch/aot_autograd.py
+++ b/torch/_functorch/aot_autograd.py
@@ -511,6 +511,8 @@ def create_aot_dispatcher_function(
                 )(*fake_flat_args)

                 req_subclass_dispatch = requires_subclass_dispatch(fake_flat_args, fw_metadata)
+                from ._aot_autograd.subclass_utils import is_dtensor_subclass_dispatch
+                dtensor_dispatch = is_dtensor_subclass_dispatch(fake_flat_args, fw_metadata)

                 if needs_autograd and not any(x.requires_grad for x in fw_metadata.output_info):
                     # We realized that none of the outputs require grad,
@@ -568,7 +570,8 @@ Found a graph input that requires gradients, and received a mutation.
 This is currently banned in the aot_export workflow. If you need this functionality, please file a github issue.

 fw_metadata={str(fw_metadata)}""")
-            if req_subclass_dispatch:
+            # NOTE: hack: make DTensor dispatch succeed!
+            if req_subclass_dispatch and not dtensor_dispatch:
                 raise RuntimeError("""\
 aot_export is not currently supported with traceable tensor subclass.
 If you need this feature, please comment on <CREATE_ISSUE_LINK>""")
diff --git a/torch/_guards.py b/torch/_guards.py
index 69912b15313..4f00d53b88e 100644
--- a/torch/_guards.py
+++ b/torch/_guards.py
@@ -817,8 +817,16 @@ def detect_fake_mode(inputs: Any = None):

     flat_inputs = pytree.tree_leaves(inputs)
     for i, flat_input in enumerate(flat_inputs):
+        from torch.utils._python_dispatch import is_traceable_wrapper_subclass
         if isinstance(flat_input, FakeTensor):
             fake_modes.append((flat_input.fake_mode, "fake tensor input", i))
+        # enable FakeMode detection in tensor wrapper_subclass
+        elif is_traceable_wrapper_subclass(flat_input):
+            attrs, _ = flat_input.__tensor_flatten__()
+            for attr in attrs:
+                inner_tensor = getattr(flat_input, attr)
+                if isinstance(inner_tensor, FakeTensor):
+                    fake_modes.append((inner_tensor.fake_mode, "fake inner tensor input", i))

     if fake_modes:
         fake_mode, desc1, i1 = fake_modes[0]
diff --git a/torch/_tensor.py b/torch/_tensor.py
index 3aa0cee639d..dd76e76e841 100644
--- a/torch/_tensor.py
+++ b/torch/_tensor.py
@@ -107,6 +107,7 @@ class Tensor(torch._C.TensorBase):
                     and self.device.type == torch._C._get_privateuse1_backend_name()
                 )
                 or (type(self) is not Tensor and self.data_ptr() == 0)
+                or type(self).__name__ == "DTensor"
             ):
                 new_tensor = self.clone()
                 if type(new_tensor) is not type(self):
diff --git a/torch/csrc/autograd/python_variable.cpp b/torch/csrc/autograd/python_variable.cpp
index ba0e913896d..0335434fbe5 100644
--- a/torch/csrc/autograd/python_variable.cpp
+++ b/torch/csrc/autograd/python_variable.cpp
@@ -656,9 +656,9 @@ static PyObject* THPVariable_make_wrapper_subclass(
       "SymInt? storage_offset=None, MemoryFormat? memory_format=None, ScalarType dtype=None, "
       "Layout layout=torch.strided, Device device=None, bool pin_memory=False, bool requires_grad=False, "
       "c10::string_view? dispatch_sizes_strides_policy=None, bool dispatch_device=False, bool dispatch_layout=False, "
-      "DispatchKeySet _extra_dispatch_keys=None)",
+      "DispatchKeySet _extra_dispatch_keys=None, SymInt? data_ptr= None)",
   });
-  ParsedArgs<14> parsed_args{};
+  ParsedArgs<15> parsed_args{};
   auto r = parser.parse(args, kwargs, parsed_args);
   PyObject* cls = r.pyobject(0);

@@ -726,8 +726,15 @@ static PyObject* THPVariable_make_wrapper_subclass(
         size_bytes,
         /*allocator=*/c10::GetAllocator(c10::kMeta),
         /*resizable=*/true};
-    // TODO: constructor should probably accept data pointer
-    storage.set_data_ptr_noswap(at::DataPtr{nullptr, r.device(7)});
+    auto data_ptr = r.toSymIntOptional(14);
+    if (data_ptr.value_or(0) != 0) {
+      // NOLINTNEXTLINE(performance-no-int-to-ptr)
+      void* p = reinterpret_cast<void*>(
+          static_cast<uintptr_t>(data_ptr->expect_int()));
+      storage.set_data_ptr_noswap(at::DataPtr{p, r.device(7)});
+    } else {
+      storage.set_data_ptr_noswap(at::DataPtr{nullptr, r.device(7)});
+    }

     auto keys = c10::DispatchKeySet({options.computeDispatchKey()});
     if (auto mb_extra_keys = r.toDispatchKeySetOptional(13)) {
@@ -2210,4 +2217,4 @@ bool THPVariable_initModule(PyObject* module) {
   torch::autograd::initTensorImplConversion(module);
   torch::utils::validate_numpy_for_dlpack_deleter_bug();
   return true;
-}
+}
\ No newline at end of file
diff --git a/torch/distributed/_functional_collectives.py b/torch/distributed/_functional_collectives.py
index a0e02292cfe..f76fded484e 100644
--- a/torch/distributed/_functional_collectives.py
+++ b/torch/distributed/_functional_collectives.py
@@ -128,6 +128,62 @@ def wait_tensor(tensor):
     return torch.ops.c10d_functional.wait_tensor(tensor)  # type: ignore[attr-defined]


+def send(self: torch.Tensor, dst: int, group: RANK_TYPES, tag: str = ""):
+    """
+    Sends the tensor to the destination process.
+
+    Args:
+        dst (int): Destination rank
+        group (ProcessGroup or List[int]): The process group to work on.
+        tag (str, optional): A unique identifier for the collective. Default: empty string
+    """
+    tag, rankset, group_size = _expand_group(group, tag)
+    tensor = torch.ops.c10d_functional.send(self, dst, tag, rankset, group_size)
+    return tensor
+
+
+def recv(self: torch.Tensor, src: int, group: RANK_TYPES, tag: str = ""):
+    """
+    Receives the tensor from the source process.
+
+    Args:
+        srd (int): Source rank
+        group (ProcessGroup or List[int]): The process group to work on.
+        tag (str, optional): A unique identifier for the collective. Default: empty string
+    """
+    tag, rankset, group_size = _expand_group(group, tag)
+    tensor = torch.ops.c10d_functional.recv(self, src, tag, rankset, group_size)
+    return tensor
+
+
+def isend(self: torch.Tensor, dst: int, group: RANK_TYPES, tag: str = ""):
+    """
+    Sends the tensor to the destination process asynchronously.
+
+    Args:
+        dst (int): Destination rank
+        group (ProcessGroup or List[int]): The process group to work on.
+        tag (str, optional): A unique identifier for the collective. Default: empty string
+    """
+    tag, rankset, group_size = _expand_group(group, tag)
+    tensor = torch.ops.c10d_functional.isend(self, dst, tag, rankset, group_size)
+    return _maybe_wrap_tensor(tensor)
+
+
+def irecv(self: torch.Tensor, src: int, group: RANK_TYPES, tag: str = ""):
+    """
+    Receives the tensor from the source process asynchronously.
+
+    Args:
+        src (int): Source rank
+        group (ProcessGroup or List[int]): The process group to work on.
+        tag (str, optional): A unique identifier for the collective. Default: empty string
+    """
+    tag, rankset, group_size = _expand_group(group, tag)
+    tensor = torch.ops.c10d_functional.irecv(self, src, tag, rankset, group_size)
+    return _maybe_wrap_tensor(tensor)
+
+
 def broadcast(self: torch.Tensor, src: int, group: RANK_TYPES, tag: str = ""):
     """
     Broadcasts the tensor to all processes in the given process group.
@@ -542,6 +598,23 @@ def _all_gather_into_tensor_coalesced_meta(self, tag, rankset, group_size):

     return [mk_out_tensor(t) for t in self]

+
+def _send_meta(self, *args):
+    return torch.empty_like(self)
+
+
+def _recv_meta(self, *args):
+    return torch.empty_like(self)
+
+
+def _isend_meta(self, *args):
+    return torch.empty_like(self)
+
+
+def _irecv_meta(self, *args):
+    return torch.empty_like(self)
+
+
 # We now register meta kernels to deal with tracing
 def _broadcast_meta(self, *args):
     return torch.empty_like(self)
@@ -619,6 +692,10 @@ def _reduce_scatter_tensor_coalesced_native_meta(inputs, reduce_op, group_size,

 def _register_ops():
     ops_defs = [
+        "send(Tensor self, int dst, str tag, int[] ranks, int group_size) -> Tensor",
+        "recv(Tensor self, int src, str tag, int[] ranks, int group_size) -> Tensor",
+        "isend(Tensor self, int dst, str tag, int[] ranks, int group_size) -> Tensor",
+        "irecv(Tensor self, int src, str tag, int[] ranks, int group_size) -> Tensor",
         "broadcast(Tensor self, int src, str tag, int[] ranks, int group_size) -> Tensor",
         "all_reduce(Tensor self, str reduceOp, str tag, int[] ranks, int group_size) -> Tensor",
         "all_reduce_coalesced(Tensor[] self, str reduceOp, str tag, int[] ranks, int group_size) -> Tensor[]",
diff --git a/torch/distributed/_functional_collectives_impl.py b/torch/distributed/_functional_collectives_impl.py
index f14ad5b067e..04445656e75 100644
--- a/torch/distributed/_functional_collectives_impl.py
+++ b/torch/distributed/_functional_collectives_impl.py
@@ -138,6 +138,37 @@ def _str_to_reduce_op(reduceOp: str) -> dist.ReduceOp:
         raise ValueError(f"Invalid reduce operation {reduceOp}")
     return cast(dist.ReduceOp, op)

+def _send(self, dst, tag, ranks, group_size):
+    group = c10d._find_or_create_pg_by_ranks_and_tag(tag, ranks, group_size)
+    assert group is not None
+
+    dist.send(self, dst, group=group)
+    return self
+
+def _recv(self, src, tag, ranks, group_size):
+    group = c10d._find_or_create_pg_by_ranks_and_tag(tag, ranks, group_size)
+    assert group is not None
+
+    dist.recv(self, src, group=group)
+    return self
+
+def _isend(self, dst, tag, ranks, group_size):
+    group = c10d._find_or_create_pg_by_ranks_and_tag(tag, ranks, group_size)
+    assert group is not None
+
+    work = dist.isend(self, dst, group=group)
+    _register_tensor_work(self, work)
+    return self
+
+
+def _irecv(self, src, tag, ranks, group_size):
+    group = c10d._find_or_create_pg_by_ranks_and_tag(tag, ranks, group_size)
+    assert group is not None
+
+    work = dist.irecv(self, src, group=group)
+    _register_tensor_work(self, work)
+    return self
+

 """
 Kernel implementations (for eager runtime only) - should never be traced by torch.compile
diff --git a/torch/distributed/_tensor/api.py b/torch/distributed/_tensor/api.py
index 068bc8b9af8..5a577046244 100644
--- a/torch/distributed/_tensor/api.py
+++ b/torch/distributed/_tensor/api.py
@@ -233,6 +233,7 @@ class DTensor(torch.Tensor):  # pyre-ignore[13]: pyre is bad at __new__
             device=local_tensor.device,
             layout=local_tensor.layout,
             requires_grad=requires_grad,
+            data_ptr=local_tensor.data_ptr(),
         )

         tensor_meta = TensorMeta(shape, stride, dtype)