Gradcheck small fixes (pytorch#53916)

Summary:
Pull Request resolved: pytorch#53916

This PR fixes some bugs that are made more clear by the previous refactor.
- make sure gradcheck returns false when its supposed to fail and when raise_exception=False.
- make sure when test_batched_grad fails, it returns false when raise_exception=False

Removing checkIfNumericalAnalyticAreClose made sense here to me because underneath its really doing `torch.allclose`, and using that directly instead of adding another opaque function to call seemed to make the code more clear.

TODO:
- ~add a test to see if when torch.allclose fails, we indeed return false.~
- ~uncomment test from previous PR.~

Test Plan: Imported from OSS

Reviewed By: heitorschueroff

Differential Revision: D27201692

Pulled By: soulitzer

fbshipit-source-id: 8b8dc37c59edb7eebc2e8db6f8839ce98a81d78b

Author: soulitzer

test/test_autograd.py

@@ -4094,6 +4094,8 @@ def test_gradcheck_check_batched_grad(self):
         # runtime error while compute batched grad (print big error)
         with self.assertRaisesRegex(RuntimeError, 'gradcheck or gradgradcheck failed while testing batched gradient'):
             gradcheck(lambda x: x.to_dense(), (x,), check_sparse_nnz=True, check_batched_grad=True)
+        self.assertFalse(gradcheck(lambda x: x.to_dense(), (x,), check_sparse_nnz=True, check_batched_grad=True,
+                                   raise_exception=False))
 
     def test_gradcheck_backward_mul_by_grad_output(self):
         # when grad_input is sparse and has incorrect sparse_dim/dense_dim
@@ -4161,6 +4163,38 @@ def hook(x):
                 gradcheck(fn, (x,))
             self.assertFalse(gradcheck(fn, (x,), raise_exception=False))
 
+    def test_gradcheck_jacobian_mismatch(self):
+        def fn(x):  # R -> R, C -> C
+            y = x.clone()
+            y.register_hook(lambda x: x + 1e-2)
+            return y
+        x = torch.ones(2, 2, requires_grad=True)
+        with self.assertRaisesRegex(RuntimeError, 'Jacobian mismatch for output 0 with respect to input 0'):
+            gradcheck(fn, (x,))
+        self.assertFalse(gradcheck(fn, (x,), raise_exception=False))
+
+        x_c = torch.ones(2, 2, requires_grad=True, dtype=torch.complex128)
+        with self.assertRaisesRegex(RuntimeError, 'Gradients failed to compare equal for grad output = 1j'):
+            gradcheck(fn, (x_c,))
+        self.assertFalse(gradcheck(fn, (x_c,), raise_exception=False))
+
+        def fn2(x):  # R -> C
+            y = torch.complex(x, x)
+            y.register_hook(lambda x: x + 1e-2)
+            return y
+        x = torch.ones(2, 2, requires_grad=True)
+        with self.assertRaisesRegex(RuntimeError, 'Gradients failed to compare equal for grad output = 1j'):
+            gradcheck(fn2, (x,))
+        self.assertFalse(gradcheck(fn2, (x,), raise_exception=False))
+
+        def fn3(x):  # C -> R
+            y = torch.real(x)
+            y.register_hook(lambda x: x + 1e-2)
+            return y
+        with self.assertRaisesRegex(RuntimeError, 'Gradients failed to compare equal for grad output = 1'):
+            gradcheck(fn3, (x_c,))
+        self.assertFalse(gradcheck(fn3, (x_c,), raise_exception=False))
+
     def test_version_counter(self):
         x = torch.randn(1, 2)
 

test/test_overrides.py

@@ -804,7 +804,6 @@ def test_gradcheck(self):
         # Tensor-likes.
         self.assertEqual(total_used_attrs, {
             'data',
-            'device',
             'dtype',
             'is_complex',
             'is_floating_point',

torch/autograd/gradcheck.py

@@ -318,7 +318,7 @@ def get_failed_batched_grad_test_msg(output_idx, input_idx, res, exp):
 """.strip()
 
 
-def test_batched_grad(fail_test, input, output, output_idx):
+def test_batched_grad(fail_test, input, output, output_idx) -> bool:
     # NB: test_batched_grad compares two autograd.grad invocations with a single
     # vmap(autograd.grad) invocation. It's not exactly a "gradcheck" in the
     # sense that we're not comparing an analytical jacobian with a numeric one,
@@ -358,6 +358,7 @@ def vjp(v):
         if torch.allclose(res, exp):
             continue
         return fail_test(get_failed_batched_grad_test_msg(output_idx, input_idx, res, exp))
+    return True
 
 
 def test_backward_mul_by_grad_output(fail_test, outputs, inputs, check_sparse_nnz) -> bool:
@@ -458,6 +459,11 @@ def _differentiable_outputs(x):
     return tuple(o for o in _as_tuple(x) if o.requires_grad)
 
 
+def get_notallclose_msg(analytical, numerical, output_idx, input_idx, error_str='') -> str:
+    return error_str + 'Jacobian mismatch for output %d with respect to input %d,\n' \
+        'numerical:%s\nanalytical:%s\n' % (output_idx, input_idx, numerical, analytical)
+
+
 # Note [VarArg of Tensors]
 # ~~~~~~~~~~~~~~~~~~~~~~~~
 # 'func' accepts a vararg of tensors, which isn't expressable in the type system at the moment.
@@ -562,31 +568,26 @@ def fn(input):
                 return False
             numerical_from_imag_grad_out = get_numerical_jacobian(fn, tupled_inputs, eps=eps, grad_out=1j)
 
-        def checkIfNumericalAnalyticAreClose(a, n, j, error_str=''):
-            if not torch.allclose(a, n, rtol, atol):
-                return fail_test(error_str + 'Jacobian mismatch for output %d with respect to input %d,\n'
-                                 'numerical:%s\nanalytical:%s\n' % (i, j, n, a))
-
         inp_tensors = iter_tensors(tupled_inputs, True)
 
         for j, (a, n, inp) in enumerate(zip(analytical, numerical, inp_tensors)):
             if a.numel() != 0 or n.numel() != 0:
-                if o.is_complex():
-                    # C -> C, R -> C
-                    a_imag_grad_out = analytical_from_imag_grad_out[j]
-                    n_imag_grad_out = numerical_from_imag_grad_out[j]
-                    checkIfNumericalAnalyticAreClose(a_imag_grad_out, n_imag_grad_out, j,
-                                                     "Gradients failed to compare equal for grad output = 1j. ")
-                if inp.is_complex():
-                    # C -> R, C -> C
-                    checkIfNumericalAnalyticAreClose(a, n, j,
-                                                     "Gradients failed to compare equal for grad output = 1. ")
-                else:
-                    # R -> R, R -> C
-                    checkIfNumericalAnalyticAreClose(a, n, j)
+                if o.is_complex():    # C -> C, R -> C
+                    if not torch.allclose(analytical_from_imag_grad_out[j], numerical_from_imag_grad_out[j], rtol, atol):
+                        return fail_test(get_notallclose_msg(analytical_from_imag_grad_out[j],
+                                                             numerical_from_imag_grad_out[j], i, j,
+                                                             "Gradients failed to compare equal for grad output = 1j. "))
+                if inp.is_complex():  # C -> R, C -> C
+                    if not torch.allclose(a, n, rtol, atol):
+                        return fail_test(get_notallclose_msg(a, n, i, j,
+                                                             "Gradients failed to compare equal for grad output = 1. "))
+                else:                 # R -> R, R -> C
+                    if not torch.allclose(a, n, rtol, atol):
+                        return fail_test(get_notallclose_msg(a, n, i, j))
 
         if check_batched_grad:
-            test_batched_grad(fail_test, tupled_inputs, o, j)
+            if not test_batched_grad(fail_test, tupled_inputs, o, i):
+                return False
 
     if not test_backward_mul_by_grad_output(fail_test, outputs, tupled_inputs, check_sparse_nnz):
         return False