[Dynamo] Support out variants of ops mutate the tensors out of the function frame (pytorch#93177)

Fixes pytorch#93136

Pull Request resolved: pytorch#93177
Approved by: https://github.com/jansel

Author: yanboliang

test/dynamo/test_dynamic_shapes.py

@@ -55,6 +55,11 @@ def make_dynamic_cls(cls):
     # Cannot call sizes() on tensor with symbolic sizes/strides
 )
 
+unittest.expectedFailure(
+    DynamicShapesReproTests.test_sort_out2_dynamic_shapes
+    # Cannot call sizes() on tensor with symbolic sizes/strides
+)
+
 # DynamicShapesExportTests
 unittest.expectedFailure(
     DynamicShapesExportTests.test_export_with_constant_list_nonzero_dynamic_shapes

test/dynamo/test_repros.py

@@ -1531,6 +1531,24 @@ def fn():
         opt_fn = torch._dynamo.optimize("eager")(fn)
         opt_fn()
 
+    def test_sort_out2(self):
+        class MyModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("sorted", torch.ones(4, 4))
+                self.register_buffer("indices", torch.ones(4, 4, dtype=torch.long))
+
+            def forward(self, x):
+                torch.sort(x, out=(self.sorted, self.indices))
+                return (x + 1, self.sorted, self.indices)
+
+        x = torch.randn(4, 4)
+        m = MyModule()
+        ref = m(x)
+        opt_m = torch._dynamo.optimize("eager")(m)
+        res = opt_m(x)
+        self.assertTrue(same(ref, res))
+
     def test_sigmoid_out(self):
 
         dtype = torch.float32
@@ -1546,6 +1564,23 @@ def fn():
         opt_fn = torch._dynamo.optimize("eager")(fn)
         opt_fn()
 
+    def test_sigmoid_out2(self):
+        class MyModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("base", torch.ones(4, 4))
+
+            def forward(self, x):
+                torch.sigmoid(x, out=self.base)
+                return x + self.base
+
+        x = torch.randn(4, 4)
+        m = MyModule()
+        ref = m(x)
+        opt_m = torch._dynamo.optimize("eager")(m)
+        res = opt_m(x)
+        self.assertTrue(same(ref, res))
+
     def test_slice_into_list_mutable(self):
         class Mod(torch.nn.Module):
             def forward(self, listy):

torch/_dynamo/variables/torch.py

@@ -491,13 +491,13 @@ def get_state_from_generator():
                         tx.find_symbolic_locals_name(x) for x in kwargs["out"].items
                     ]
                     for idx, name in enumerate(output_tensor_names):
-                        assert name in tx.symbolic_locals
-                        tx.symbolic_locals[name] = tensor_variable.items[idx]
+                        if name in tx.symbolic_locals:
+                            tx.symbolic_locals[name] = tensor_variable.items[idx]
                 elif isinstance(tensor_variable, TensorVariable):
                     assert isinstance(kwargs["out"], TensorVariable)
                     name = tx.find_symbolic_locals_name(kwargs["out"])
-                    assert name in tx.symbolic_locals
-                    tx.symbolic_locals[name] = tensor_variable
+                    if name in tx.symbolic_locals:
+                        tx.symbolic_locals[name] = tensor_variable
                 else:
                     unimplemented(f"out variant of {type(kwargs['out'])}")