Fix grad higher order handling TupleVariable (pytorch#106425)

Previously, we assume the argnums is a **ConstantVariable**. However I accidentally triggered an error on CI where argnums could be a **TupleVariable**. In that case, we have an attribute error when access the .value of argnums.

This PR adds support for the TupleVariable. It allows the unit test to pass without falling back to eager
"PYTORCH_TEST_WITH_DYNAMO=1 python test/functorch/test_eager_transforms.py -k test_argnums_cpu"

Test Plan:
see modified test.

Pull Request resolved: pytorch#106425
Approved by: https://github.com/yanboliang, https://github.com/anijain2305, https://github.com/kshitij12345

Author: ydwu4

test/dynamo/test_higher_order_ops.py

@@ -1851,15 +1851,21 @@ def forward(self, l_x_, l_y_):
     def test_grad_two_tensor_all_grad_has_aux(self):
         counters.clear()
 
+        nums = (0, 1)
+
         def fn(x, y):
             return ((x.sin() + y).sum(), x.cos())
 
-        def wrapper_fn(x, y):
+        def wrapper_fn_const_var(x, y):
             return torch.func.grad(fn, argnums=(0, 1), has_aux=True)(x, y)
 
+        def wrapper_fn_tuple_var(x, y):
+            return torch.func.grad(fn, argnums=nums, has_aux=True)(x, y)
+
         y = torch.randn(3, 3, 3)
         x = torch.randn(3, 3, 3)
-        wrapped_gm = self._grad_compile_check(wrapper_fn, (x, y))
+        wrapped_gm_const_var = self._grad_compile_check(wrapper_fn_const_var, (x, y))
+        wrapped_gm_tuple_var = self._grad_compile_check(wrapper_fn_tuple_var, (x, y))
 
         # Dynamic shapes produce a slightly different graph.
         if check_dynamic_shape_capture():
@@ -1894,8 +1900,14 @@ def forward(self, l_x_, l_y_):
             _set_grad_enabled_1 = torch._C._set_grad_enabled(True)
             return (sum_1, cos)
 """
-        actual = normalize_gm(wrapped_gm.print_readable(print_output=False))
-        self.assertExpectedInline(actual, expected)
+        actual_const_var = normalize_gm(
+            wrapped_gm_const_var.print_readable(print_output=False)
+        )
+        actual_tuple_var = normalize_gm(
+            wrapped_gm_tuple_var.print_readable(print_output=False)
+        )
+        self.assertExpectedInline(actual_const_var, expected)
+        self.assertExpectedInline(actual_tuple_var, expected)
 
     def test_grad_over_grad(self):
         counters.clear()

torch/_dynamo/variables/higher_order_ops.py

@@ -666,15 +666,37 @@ def call_function(
         # For has_aux=False, Tuple[gradients of inputs indicated by argnums].
         # For has_aux=True, Tuple[Tuple[gradients of inputs indicated by argnums], aux values]
         # NOTE: example_value should match `grad_output`.
-        if isinstance(argnums.value, int):
-            example_value = (
-                args[argnums.value].as_proxy().node.meta["example_value"].contiguous()
-            )
-        else:
-            example_value = tuple(
-                args[idx].as_proxy().node.meta["example_value"].contiguous()
-                for idx in argnums.value
-            )
+        def _from_args(idx):
+            return args[idx].as_proxy().node.meta["example_value"].contiguous()
+
+        def to_python_ints(argnums):
+            if not isinstance(argnums, (ConstantVariable, TupleVariable)):
+                raise UserError(
+                    UserErrorType.INVALID_INPUT,
+                    f"argnums is expected to be int or tuple of ints. Got {argnums}.",
+                )
+
+            if isinstance(argnums, ConstantVariable):
+                if not isinstance(argnums.value, (int, tuple)):
+                    raise UserError(
+                        UserErrorType.INVALID_INPUT,
+                        f"argnums is expected to be int or tuple of ints. Got {argnums}.",
+                    )
+                return argnums.value
+            else:
+                const_vars = argnums.unpack_var_sequence(tx)
+                if not all(
+                    isinstance(var, ConstantVariable) and isinstance(var.value, int)
+                    for var in const_vars
+                ):
+                    raise UserError(
+                        UserErrorType.INVALID_INPUT,
+                        f"argnums is expected to contain int only. Got {const_vars}.",
+                    )
+                return tuple(var.value for var in const_vars)
+
+        argnums_v = to_python_ints(argnums)
+        example_value = pytree.tree_map(_from_args, argnums_v)
 
         if has_aux.value:
             # case : has_aux = True
@@ -691,12 +713,12 @@ def call_function(
 
         # Call contiguous on all the computed grads.
         if not has_aux.value:
-            if isinstance(argnums.value, int):
+            if isinstance(argnums_v, int):
                 return fx_proxy.call_method(tx, "contiguous", (), {})
             else:
                 grads = fx_proxy
                 items = []
-                for idx in range(len(argnums.value)):
+                for idx in range(len(argnums_v)):
                     proxy = grads.call_method(
                         tx, "__getitem__", (ConstantVariable(idx),), {}
                     ).call_method(tx, "contiguous", (), {})
@@ -706,11 +728,11 @@ def call_function(
             # fx_proxy -> Tuple(grads, aux)
             grads = fx_proxy.call_method(tx, "__getitem__", (ConstantVariable(0),), {})
             aux = fx_proxy.call_method(tx, "__getitem__", (ConstantVariable(1),), {})
-            if isinstance(argnums.value, int):
+            if isinstance(argnums_v, int):
                 return TupleVariable([grads.call_method(tx, "contiguous", (), {}), aux])
             else:
                 items = []
-                for idx in range(len(argnums.value)):
+                for idx in range(len(argnums_v)):
                     proxy = grads.call_method(
                         tx, "__getitem__", (ConstantVariable(idx),), {}
                     ).call_method(tx, "contiguous", (), {})