Aot Autograd - Dont fallback with new op

Author: anijain2305

benchmarks/common.py

@@ -68,7 +68,6 @@ def set_model_name(name):
     "pytorch_struct",
     "speech_transformer",
     "vision_maskrcnn",
-    "moco",
     # Huggingface
     "AlbertForMaskedLM",  # OOM
     "AlbertForQuestionAnswering",  # OOM

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

@@ -34,8 +34,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):
@@ -53,11 +52,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: