Revert "make sure dynamo doesn't inline DTensor __new__ or __torch_dispatch__ (#123347)"

pytorchmergebot · pytorchmergebot · commit 9d88339b535f · 2024-04-16T22:08:24.000Z
This reverts commit 63dcb5b. Reverted #123347 on behalf of https://github.com/facebook-github-bot due to Diff reverted internally ([comment](#123347 (comment)))
diff --git a/test/distributed/_tensor/test_dtensor_compile.py b/test/distributed/_tensor/test_dtensor_compile.py
@@ -191,47 +191,6 @@ def fn(x):
         res = opt_fn(x)
         self.assertEqual(res, ref)
 
-    def test_dtensor_constructor_w_graph_break(self):
-        mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
-
-        # test passing in DTensor as inputs/outputs and run some tensor computation
-        def fn(x):
-            print("graph break!")
-            return DTensor(
-                x,
-                mesh,
-                (Replicate(), Shard(0)),
-                shape=[128, 32],
-                dtype=x.dtype,
-                requires_grad=x.requires_grad,
-                stride=[32, 1],
-            )
-
-        x = torch.randn(64, 32, requires_grad=True)
-        out = fn(x)
-        out2 = torch.compile(fn, backend="eager")(x)
-
-    def test_dtensor_constructor_w_dynamo_disable(self):
-        mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
-
-        @torch._dynamo.disable(recursive=False)
-        def fn(x):
-            print("foo")
-            return DTensor(
-                x,
-                mesh,
-                (Replicate(),),
-                shape=torch.Size([32]),
-                dtype=x.dtype,
-                requires_grad=x.requires_grad,
-                stride=(1,),
-            )
-
-        x = torch.randn(32, requires_grad=True)
-        out = fn(x)
-        out2 = torch.compile(fn, backend="eager")(x)
-        self.assertEqual(out, out2)
-
     def test_dtensor_noncontiguous_output(self):
         mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
 
diff --git a/torch/distributed/_tensor/api.py b/torch/distributed/_tensor/api.py
@@ -198,7 +198,6 @@ class DTensor(torch.Tensor):  # pyre-ignore[13]: pyre is bad at __new__
     _op_dispatcher: op_dispatch.OpDispatcher = op_dispatch.OpDispatcher()
 
     @staticmethod
-    @torch._dynamo.disable
     def __new__(
         cls,
         local_tensor: torch.Tensor,
@@ -289,7 +288,6 @@ def __coerce_same_metadata_as_tangent__(self, metadata_tensor):
         )
 
     @classmethod
-    @torch._dynamo.disable
     # pyre-fixme[3]: Return type must be annotated.
     # pyre-fixme[2]: Parameter must be annotated.
     def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
