[dynamo][hops] Add xfail tests for side effects (pytorch#168394)

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

Author: anijain2305

test/dynamo/test_autograd_function.py

@@ -2,6 +2,7 @@
 # flake8: noqa: B950
 import copy
 import math
+import unittest
 from dataclasses import dataclass
 
 import torch
@@ -1543,6 +1544,43 @@ def f(x, y):
         loss.backward()
         self.assertEqual(x + y, z)
 
+    @unittest.expectedFailure
+    def test_nonlocal_list_mutation_in_autograd_function(self):
+        """Test that nonlocal list mutation in autograd.Function forward is handled correctly."""
+
+        class SimpleAutogradFunc(torch.autograd.Function):
+            @staticmethod
+            def forward(ctx, x, z):
+                # Simple computation
+                o = torch.matmul(x, x) @ x
+                out = x.sin()
+                # Mutate the nonlocal list
+                z.append(out)
+                return torch.cos(torch.sin(o)), torch.sin(x)
+
+            @staticmethod
+            def backward(ctx, grad_output1, grad_output2):
+                # Simple backward
+                return grad_output1 + grad_output2, None
+
+        def fn(x):
+            z = []
+
+            outs = SimpleAutogradFunc.apply(x, z)
+            out1 = outs[0]
+            # Check that the extra output pytree handling is done properly
+            out2 = outs[-1]
+
+            return out1 + out2, z[0]
+
+        x = torch.randn(4, 4, requires_grad=True)
+        ref = fn(x)
+
+        opt_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+        res = opt_fn(x)
+        self.assertEqual(ref[0], res[0])
+        self.assertEqual(ref[1], res[1])
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests

test/higher_order_ops/test_invoke_subgraph.py

@@ -910,6 +910,35 @@ def forward(self, a: "f32[8]", l_y_: "f32[8]"):
 """,
             )
 
+    @unittest.expectedFailure
+    def test_nonlocal_list_mutation_hidden(self):
+        """Test that nonlocal list mutation inside nested_compile_region is handled correctly."""
+
+        @nested_compile_region
+        def gn(x, z):
+            o = torch.matmul(x, x) @ x
+            out = x.sin()
+            z.append(out)
+            return torch.cos(torch.sin(o)), torch.sin(x)
+
+        def fn(x):
+            z = []
+
+            outs = gn(x, z)
+            out1 = outs[0]
+            # Check that the extra output pytree handling is done properly
+            out2 = outs[-1]
+
+            return out1 + out2, z[0]
+
+        x = torch.randn(4, 4, requires_grad=True)
+        ref = fn(x)
+
+        opt_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
+        res = opt_fn(x)
+        self.assertEqual(ref[0], res[0])
+        self.assertEqual(ref[1], res[1])
+
     @inductor_config.patch("fx_graph_cache", False)
     def test_view_to_reshape(self):
         @nested_compile_region