[fix] jacrev and jacfwd : support non-tensor args again (pytorch#97746)

Fixes pytorch#97636

The code to check if argument tensor are complex assumed that all arguments are tensor (which is not the case) which lead to the error.

Pull Request resolved: pytorch#97746
Approved by: https://github.com/zou3519

test/functorch/test_eager_transforms.py

@@ -2216,6 +2216,16 @@ def fn(x):
         with self.assertRaisesRegex(RuntimeError, "jacfwd: Expected all outputs"):
             jacfwd(fn)(x)
 
+    @jacrev_and_jacfwd
+    def test_jac_with_non_tensor_args(self, device, jacapi):
+        def f(t, int_x):
+            return t + int_x
+
+        t = torch.randn(3, 3, device=device)
+
+        actual = jacapi(f)(t, 3)
+        expected = torch.autograd.functional.jacobian(partial(f, int_x=3), t)
+        self.assertEqual(actual, expected)
 
 class TestHessian(TestCase):
     def _test_against_reference(self, f, inputs):

torch/_functorch/eager_transforms.py

@@ -344,7 +344,7 @@ def _safe_zero_index(x):
 def error_if_complex(func_name, args, is_input):
     flat_args, _ = tree_flatten(args)
     for idx, arg in enumerate(flat_args):
-        if arg.dtype.is_complex:
+        if isinstance(arg, torch.Tensor) and arg.dtype.is_complex:
             input_or_output = ("inputs" if is_input else "outputs")
             err_msg = (f"{func_name}: Expected all {input_or_output} "
                        f"to be real but received complex tensor at flattened input idx: {idx}")