[dynamo] Support Tuple[int] args to autograd.Function (#123887)

Pull Request resolved: #123887
Approved by: https://github.com/anijain2305
ghstack dependencies: #123700, #123705, #123786, #123790, #123803, #123804, #123896

Author: jansel

test/dynamo/test_autograd_function.py

@@ -924,6 +924,33 @@ def foo(x):
         foo(torch.randn(2))
         foo(torch.randn(2, requires_grad=True))
 
+    def test_tuple_arg(self):
+        cnt = torch._dynamo.testing.CompileCounter()
+
+        class TupleArgFunc(torch.autograd.Function):
+            @staticmethod
+            def forward(ctx, x, shape):
+                ctx.save_for_backward(torch.randn(shape))
+                return x + 1
+
+            @staticmethod
+            def backward(ctx, grad_output):
+                (result,) = ctx.saved_tensors
+                return result, None
+
+        @torch.compile(backend=cnt, fullgraph=True)
+        def fn():
+            return TupleArgFunc.apply(x, shape)
+
+        shape = (10, 10)
+        x = torch.randn(shape, requires_grad=True)
+        out = fn()
+        out.sum().backward()
+        self.assertEqual(out, x + 1)
+        self.assertEqual(x.grad.shape, shape)
+        self.assertEqual(cnt.frame_count, 1)
+        self.assertEqual(cnt.op_count, 2)
+
     @requires_cuda
     @skipIfRocm
     def test_triton_kernel_basic(self):

torch/_dynamo/variables/higher_order_ops.py

@@ -137,7 +137,7 @@ def validate_args_and_maybe_create_graph_inputs(
     set_subgraph_inputs,
     description,
 ):
-    from . import AutogradFunctionContextVariable, EnumVariable
+    from . import AutogradFunctionContextVariable
     from .builder import wrap_fx_proxy_cls
 
     assert tracer.parent is not None
@@ -166,7 +166,7 @@ def validate_args_and_maybe_create_graph_inputs(
                 args.append(a)
                 continue
 
-            if isinstance(a, (ConstantVariable, EnumVariable)):
+            if a.is_python_constant():
                 # This arg is not used in the body of the higher order op.
                 # Currently, this new input is added to make the calls
                 # happy, which expect a fixed number of arguments. In