Aot Autograd - Dont fallback with new op (#1493)

* Aot Autograd - Dont fallback with new op

* Bring back moco - Moco fails with einsum float16 error

* CI

Author: anijain2305

test/test_aot_autograd.py

@@ -60,23 +60,6 @@ def fn(param, y):
         aot_fn(x, y)
         self.assertTrue(not is_safe[0])
 
-    def test_new(self):
-        # https://github.com/pytorch/torchdynamo/issues/1448
-        def fn(argsort: torch.Tensor):
-            new = argsort.new(2, 12, 4096)
-            x = torch.add(new, 2)
-            return (
-                new,
-                x,
-            )
-
-        x = torch.randn((2, 12, 4096))
-        is_safe = [True]
-        compiler_fn = functools.partial(compiler_safe_fn, is_safe=is_safe)
-        aot_fn = torchdynamo.optimize(compiler_fn)(fn)
-        aot_fn(x)
-        self.assertTrue(not is_safe[0])
-
     def test_negative_testing(self):
         def fn(x, y):
             return torch.sin(x).add_(y)

torchdynamo/optimizations/training.py

@@ -33,8 +33,7 @@ def is_aot_autograd_safe_to_run(gm, example_inputs):
     Issues
     1) LSTM - https://github.com/pytorch/torchdynamo/issues/1147
     2) LSTM - https://github.com/pytorch/functorch/issues/586
-    3) New op - https://github.com/pytorch/torchdynamo/issues/1448
-    4) Input mutation - https://github.com/pytorch/torchdynamo/issues/1301
+    3) Input mutation - https://github.com/pytorch/torchdynamo/issues/1301
     """
 
     def raise_or_warn(reason):
@@ -52,11 +51,6 @@ def raise_or_warn(reason):
         if submod.__class__.__name__ == "LSTM":
             return raise_or_warn("LSTM")
 
-    # 2) new does not work with fake tensor and aot autograd
-    for node in gm.graph.nodes:
-        if node.op == "call_method" and node.target == "new":
-            return raise_or_warn("new operator")
-
     # 2) Mutation in the graph
     mutated = False
     try: