Improve error messages of torch.testing.assert_close in case of mismatching values (#60091)

pmeier · facebook-github-bot · commit 9979289037f5 · 2021-07-07T09:54:09.000-07:00
Summary: Pull Request resolved: #60091 Closes #58383. (1) and (2) are implemented. (3) was rejected. No consensus was reached on (4) and (5). Improvements: - Instead of calling everything "Tensors" we now use "Scalars" and "Tensor-likes" depending on the shape. Plus, we now internally have the option to adapt this identifier for example to report "Imaginary components of complex tensor-likes", which is even more expressive. - The reported conditions "not close" and "not equal" are now determined based on `rtol` and `atol`. - The number of mismatched elements and the offending indices are only reported in case the inputs are not scalar - The allowed `rtol` and `atol` is only reported if `> 0` **Example 1** ```python torch.testing.assert_close(1, 3, rtol=0, atol=1) ``` Before: ``` AssertionError: Tensors are not close! Mismatched elements: 1 / 1 (100.0%) Greatest absolute difference: 2 at 0 (up to 1 allowed) Greatest relative difference: 0.6666666865348816 at 0 (up to 0 allowed) ``` After: ``` AssertionError: Scalars are not close! Absolute difference: 2 (up to 1 allowed) Relative difference: 0.6666666865348816 ``` **Example 2** ```python torch.manual_seed(0) t = torch.rand((2, 2), dtype=torch.complex64) torch.testing.assert_close(t, t + complex(0, 1)) ``` Before: ``` AssertionError: Tensors are not close! Mismatched elements: 4 / 4 (100.0%) Greatest absolute difference: 1.0000000596046448 at (0, 0) (up to 1e-05 allowed) Greatest relative difference: 0.8833684352411922 at (0, 1) (up to 1.3e-06 allowed) The failure occurred for the imaginary part. ``` After: ``` AssertionError: Imaginary components of tensor-likes are not close! Mismatched elements: 4 / 4 (100.0%) Greatest absolute difference: 1.0000000596046448 at index (0, 0) (up to 1e-05 allowed) Greatest relative difference: 0.8833684352411922 at index (0, 1) (up to 1.3e-06 allowed) ``` Test Plan: Imported from OSS Reviewed By: ngimel Differential Revision: D29556357 Pulled By: mruberry fbshipit-source-id: 559d4a19ad4fc069b2b4f8cb5fc2f6058621e33d
diff --git a/test/test_testing.py b/test/test_testing.py
@@ -915,6 +915,39 @@ def test_mismatching_device_no_check(self, devices):
 
 
 class TestAssertCloseErrorMessage(TestCase):
+    def test_identifier_tensor_likes(self):
+        actual = torch.tensor([1, 2, 3, 4])
+        expected = torch.tensor([1, 2, 5, 6])
+
+        for fn in assert_close_with_inputs(actual, expected):
+            with self.assertRaisesRegex(AssertionError, re.escape("Tensor-likes")):
+                fn()
+
+    def test_identifier_scalars(self):
+        actual = 3
+        expected = 5
+        for fn in assert_close_with_inputs(actual, expected):
+            with self.assertRaisesRegex(AssertionError, re.escape("Scalars")):
+                fn()
+
+    def test_not_equal(self):
+        actual = torch.tensor([1, 2, 3, 4], dtype=torch.float32)
+        expected = torch.tensor([1, 2, 5, 6], dtype=torch.float32)
+
+        for fn in assert_close_with_inputs(actual, expected):
+            with self.assertRaisesRegex(AssertionError, re.escape("not equal")):
+                fn(rtol=0.0, atol=0.0)
+
+    def test_not_close(self):
+        actual = torch.tensor([1, 2, 3, 4], dtype=torch.float32)
+        expected = torch.tensor([1, 2, 5, 6], dtype=torch.float32)
+
+        for fn, (rtol, atol) in itertools.product(
+            assert_close_with_inputs(actual, expected), ((1.3e-6, 0.0), (0.0, 1e-5), (1.3e-6, 1e-5))
+        ):
+            with self.assertRaisesRegex(AssertionError, re.escape("not close")):
+                fn(rtol=rtol, atol=atol)
+
     def test_mismatched_elements(self):
         actual = torch.tensor([1, 2, 3, 4])
         expected = torch.tensor([1, 2, 5, 6])
@@ -928,15 +961,31 @@ def test_abs_diff(self):
         expected = torch.tensor([[1, 2], [5, 4]])
 
         for fn in assert_close_with_inputs(actual, expected):
-            with self.assertRaisesRegex(AssertionError, re.escape("Greatest absolute difference: 2 at (1, 0)")):
+            with self.assertRaisesRegex(AssertionError, re.escape("Greatest absolute difference: 2 at index (1, 0)")):
+                fn()
+
+    def test_abs_diff_scalar(self):
+        actual = 3
+        expected = 5
+
+        for fn in assert_close_with_inputs(actual, expected):
+            with self.assertRaisesRegex(AssertionError, re.escape("Absolute difference: 2")):
                 fn()
 
     def test_rel_diff(self):
         actual = torch.tensor([[1, 2], [3, 4]])
         expected = torch.tensor([[1, 4], [3, 4]])
 
         for fn in assert_close_with_inputs(actual, expected):
-            with self.assertRaisesRegex(AssertionError, re.escape("Greatest relative difference: 0.5 at (0, 1)")):
+            with self.assertRaisesRegex(AssertionError, re.escape("Greatest relative difference: 0.5 at index (0, 1)")):
+                fn()
+
+    def test_rel_diff_scalar(self):
+        actual = 2
+        expected = 4
+
+        for fn in assert_close_with_inputs(actual, expected):
+            with self.assertRaisesRegex(AssertionError, re.escape("Relative difference: 0.5")):
                 fn()
 
     def test_zero_div_zero(self):
@@ -952,25 +1001,21 @@ def test_zero_div_zero(self):
     def test_rtol(self):
         rtol = 1e-3
 
-        actual = torch.tensor(1)
-        expected = torch.tensor(2)
+        actual = torch.tensor((1, 2))
+        expected = torch.tensor((2, 2))
 
         for fn in assert_close_with_inputs(actual, expected):
-            with self.assertRaisesRegex(
-                AssertionError, re.escape(f"Greatest relative difference: 0.5 at 0 (up to {rtol} allowed)")
-            ):
+            with self.assertRaisesRegex(AssertionError, re.escape(f"(up to {rtol} allowed)")):
                 fn(rtol=rtol, atol=0.0)
 
     def test_atol(self):
         atol = 1e-3
 
-        actual = torch.tensor(1)
-        expected = torch.tensor(2)
+        actual = torch.tensor((1, 2))
+        expected = torch.tensor((2, 2))
 
         for fn in assert_close_with_inputs(actual, expected):
-            with self.assertRaisesRegex(
-                AssertionError, re.escape(f"Greatest absolute difference: 1 at 0 (up to {atol} allowed)")
-            ):
+            with self.assertRaisesRegex(AssertionError, re.escape(f"(up to {atol} allowed)")):
                 fn(rtol=0.0, atol=atol)
 
     def test_msg_str(self):
@@ -1049,22 +1094,6 @@ def test_mismatching_nan_with_equal_nan_relaxed(self):
         for fn in assert_close_with_inputs(actual, expected):
             fn(equal_nan="relaxed")
 
-    def test_mismatching_values_msg_real(self):
-        actual = torch.tensor(complex(0, 1))
-        expected = torch.tensor(complex(1, 1))
-
-        for fn in assert_close_with_inputs(actual, expected):
-            with self.assertRaisesRegex(AssertionError, re.escape("The failure occurred for the real part")):
-                fn()
-
-    def test_mismatching_values_msg_imag(self):
-        actual = torch.tensor(complex(1, 0))
-        expected = torch.tensor(complex(1, 1))
-
-        for fn in assert_close_with_inputs(actual, expected):
-            with self.assertRaisesRegex(AssertionError, re.escape("The failure occurred for the imaginary part")):
-                fn()
-
     def test_matching_conjugate_bit(self):
         actual = torch.tensor(complex(1, 1)).conj()
         expected = torch.tensor(complex(1, -1))
diff --git a/torch/testing/_asserts.py b/torch/testing/_asserts.py
@@ -86,11 +86,11 @@ def wrapper(
 
         error_meta = check_tensors(actual.real, expected.real, equal_nan=equal_nan, **kwargs)
         if error_meta:
-            return error_meta.amend_msg(postfix="\n\nThe failure occurred for the real part.")
+            return error_meta
 
         error_meta = check_tensors(actual.imag, expected.imag, equal_nan=equal_nan, **kwargs)
         if error_meta:
-            return error_meta.amend_msg(postfix="\n\nThe failure occurred for the imaginary part.")
+            return error_meta
 
         return None
 
@@ -132,30 +132,6 @@ def wrapper(actual: Tensor, expected: Tensor, **kwargs: Any) -> Optional[_Testin
     return wrapper
 
 
-def _check_quantized(
-    check_tensor_values: Callable[..., Optional[_TestingErrorMeta]]
-) -> Callable[..., Optional[_TestingErrorMeta]]:
-    """Decorates non-quantized tensor check functions to handle quantized tensors.
-
-    If the inputs are not quantized, this decorator is a no-op.
-
-    Args:
-        check_tensor_values (Callable[..., Optional[_TestingErrorMeta]]): Tensor check function for continuous tensors.
-
-    Returns:
-        Optional[_TestingErrorMeta]: Return value of :attr:`check_tensors`.
-    """
-
-    @functools.wraps(check_tensor_values)
-    def wrapper(actual: Tensor, expected: Tensor, **kwargs: Any) -> Optional[_TestingErrorMeta]:
-        if not actual.is_quantized:
-            return check_tensor_values(actual, expected, **kwargs)
-
-        return check_tensor_values(actual.dequantize(), expected.dequantize(), **kwargs)
-
-    return wrapper
-
-
 def _check_sparse_csr_members_individually(
     check_tensors: Callable[..., Optional[_TestingErrorMeta]]
 ) -> Callable[..., Optional[_TestingErrorMeta]]:
@@ -191,6 +167,30 @@ def wrapper(actual: Tensor, expected: Tensor, **kwargs: Any) -> Optional[_Testin
     return wrapper
 
 
+def _check_quantized(
+    check_tensor_values: Callable[..., Optional[_TestingErrorMeta]]
+) -> Callable[..., Optional[_TestingErrorMeta]]:
+    """Decorates non-quantized tensor check functions to handle quantized tensors.
+
+    If the inputs are not quantized, this decorator is a no-op.
+
+    Args:
+        check_tensor_values (Callable[..., Optional[_TestingErrorMeta]]): Tensor check function for continuous tensors.
+
+    Returns:
+        Optional[_TestingErrorMeta]: Return value of :attr:`check_tensors`.
+    """
+
+    @functools.wraps(check_tensor_values)
+    def wrapper(actual: Tensor, expected: Tensor, **kwargs: Any) -> Optional[_TestingErrorMeta]:
+        if not actual.is_quantized:
+            return check_tensor_values(actual, expected, **kwargs)
+
+        return check_tensor_values(actual.dequantize(), expected.dequantize(), **kwargs)
+
+    return wrapper
+
+
 def _check_supported_tensor(input: Tensor) -> Optional[_TestingErrorMeta]:
     """Checks if the tensor is supported by the current infrastructure.
 
@@ -298,7 +298,9 @@ def _equalize_attributes(actual: Tensor, expected: Tensor) -> Tuple[Tensor, Tens
 DiagnosticInfo = SimpleNamespace
 
 
-def _trace_mismatches(actual: Tensor, expected: Tensor, mismatches: Tensor) -> DiagnosticInfo:
+def _trace_mismatches(
+    actual: Tensor, expected: Tensor, mismatches: Tensor, *, rtol: float, atol: float
+) -> DiagnosticInfo:
     """Traces mismatches and returns diagnostic information.
 
     Args:
@@ -312,17 +314,17 @@ def _trace_mismatches(actual: Tensor, expected: Tensor, mismatches: Tensor) -> D
 
             - ``number_of_elements`` (int): Number of elements in each tensor being compared.
             - ``total_mismatches`` (int): Total number of mismatches.
-            - ``mismatch_ratio`` (float): Total mismatches divided by number of elements.
             - ``max_abs_diff`` (Union[int, float]): Greatest absolute difference of the inputs.
             - ``max_abs_diff_idx`` (Union[int, Tuple[int, ...]]): Index of greatest absolute difference.
+            - ``atol`` (float): Allowed absolute tolerance.
             - ``max_rel_diff`` (Union[int, float]): Greatest relative difference of the inputs.
             - ``max_rel_diff_idx`` (Union[int, Tuple[int, ...]]): Index of greatest relative difference.
+            - ``rtol`` (float): Allowed relative tolerance.
 
             For ``max_abs_diff`` and ``max_rel_diff`` the type depends on the :attr:`~torch.Tensor.dtype` of the inputs.
     """
     number_of_elements = mismatches.numel()
     total_mismatches = torch.sum(mismatches).item()
-    mismatch_ratio = total_mismatches / number_of_elements
 
     dtype = torch.float64 if actual.dtype.is_floating_point else torch.int64
     a_flat = actual.flatten().to(dtype)
@@ -342,14 +344,59 @@ def _trace_mismatches(actual: Tensor, expected: Tensor, mismatches: Tensor) -> D
     return SimpleNamespace(
         number_of_elements=number_of_elements,
         total_mismatches=cast(int, total_mismatches),
-        mismatch_ratio=mismatch_ratio,
         max_abs_diff=max_abs_diff.item(),
         max_abs_diff_idx=_unravel_index(max_abs_diff_flat_idx.item(), mismatches.shape),
+        atol=atol,
         max_rel_diff=max_rel_diff.item(),
         max_rel_diff_idx=_unravel_index(max_rel_diff_flat_idx.item(), mismatches.shape),
+        rtol=rtol,
     )
 
 
+def _make_mismatch_msg(
+    actual: Tensor,
+    expected: Tensor,
+    trace: DiagnosticInfo,
+    *,
+    identifier: Optional[Union[str, Callable[[str], str]]] = None,
+) -> str:
+    scalar_comparison = actual.size() == torch.Size([])
+    equality = trace.rtol == 0 and trace.atol == 0
+
+    def append_difference(msg: str, *, type: str, difference: float, index: Tuple[int, ...], tolerance: float) -> str:
+        if scalar_comparison:
+            msg += f"{type.title()} difference: {difference}"
+        else:
+            msg += f"Greatest {type} difference: {difference} at index {index}"
+        if not equality:
+            msg += f" (up to {tolerance} allowed)"
+        msg += "\n"
+        return msg
+
+    default_identifier = "Scalars" if scalar_comparison else "Tensor-likes"
+    if identifier is None:
+        identifier = default_identifier
+    elif callable(identifier):
+        identifier = identifier(default_identifier)
+
+    msg = f"{identifier} are not {'equal' if equality else 'close'}!\n\n"
+
+    if not scalar_comparison:
+        msg += (
+            f"Mismatched elements: {trace.total_mismatches} / {trace.number_of_elements} "
+            f"({trace.total_mismatches / trace.number_of_elements:.1%})\n"
+        )
+
+    msg = append_difference(
+        msg, type="absolute", difference=trace.max_abs_diff, index=trace.max_abs_diff_idx, tolerance=trace.atol
+    )
+    msg = append_difference(
+        msg, type="relative", difference=trace.max_rel_diff, index=trace.max_rel_diff_idx, tolerance=trace.rtol
+    )
+
+    return msg
+
+
 @_check_quantized
 @_check_sparse_coo_members_individually
 @_check_sparse_csr_members_individually
@@ -371,7 +418,7 @@ def _check_values_close(
         rtol (float): Relative tolerance.
         atol (float): Absolute tolerance.
         equal_nan (bool): If ``True``, two ``NaN`` values will be considered equal.
-        msg (Optional[Union[str, Callable[[Tensor, Tensor, SimpleNamespace], str]]]): Optional error message. Can be
+        msg (Optional[Union[str, Callable[[Tensor, Tensor, DiagnosticInfo], str]]]): Optional error message. Can be
             passed as callable in which case it will be called with the inputs and the result of
             :func:`_trace_mismatches`.
 
@@ -383,15 +430,10 @@ def _check_values_close(
     if not torch.any(mismatches):
         return None
 
-    trace = _trace_mismatches(actual, expected, mismatches)
+    trace = _trace_mismatches(actual, expected, mismatches, rtol=rtol, atol=atol)
     if msg is None:
-        msg = (
-            f"Tensors are not close!\n\n"
-            f"Mismatched elements: {trace.total_mismatches} / {trace.number_of_elements} ({trace.mismatch_ratio:.1%})\n"
-            f"Greatest absolute difference: {trace.max_abs_diff} at {trace.max_abs_diff_idx} (up to {atol} allowed)\n"
-            f"Greatest relative difference: {trace.max_rel_diff} at {trace.max_rel_diff_idx} (up to {rtol} allowed)"
-        )
-    elif callable(msg):
+        msg = _make_mismatch_msg
+    if callable(msg):
         msg = msg(actual, expected, trace)
     return _TestingErrorMeta(AssertionError, msg)
 
@@ -407,7 +449,7 @@ def _check_tensors_close(
     check_dtype: bool = True,
     check_stride: bool = True,
     check_is_coalesced: bool = True,
-    msg: Optional[Union[str, Callable[[Tensor, Tensor, SimpleNamespace], str]]] = None,
+    msg: Union[str, Callable[[Tensor, Tensor, DiagnosticInfo], str]],
 ) -> Optional[_TestingErrorMeta]:
     r"""Checks that the values of :attr:`actual` and :attr:`expected` are close.
 
@@ -422,7 +464,7 @@ def _check_tensors_close(
     only considered close if and only if they are equal. ``NaN``'s are only considered equal to each other if
     :attr:`equal_nan` is ``True``.
 
-    For a description of the parameters see :func:`assert_equal`.
+    For a description of the parameters see :func:`assert_close`.
 
     Returns:
         Optional[_TestingErrorMeta]: If checks did not pass.
@@ -763,11 +805,12 @@ def assert_close(
 
     - ``number_of_elements`` (int): Number of elements in each tensor being compared.
     - ``total_mismatches`` (int): Total number of mismatches.
-    - ``mismatch_ratio`` (float): Total mismatches divided by number of elements.
     - ``max_abs_diff`` (Union[int, float]): Greatest absolute difference of the inputs.
     - ``max_abs_diff_idx`` (Union[int, Tuple[int, ...]]): Index of greatest absolute difference.
+    - ``atol`` (float): Allowed absolute tolerance.
     - ``max_rel_diff`` (Union[int, float]): Greatest relative difference of the inputs.
     - ``max_rel_diff_idx`` (Union[int, Tuple[int, ...]]): Index of greatest relative difference.
+    - ``rtol`` (float): Allowed relative tolerance.
 
     For ``max_abs_diff`` and ``max_rel_diff`` the type depends on the :attr:`~torch.Tensor.dtype` of the inputs.
 
@@ -858,9 +901,10 @@ def assert_close(
         AssertionError: Argh, the tensors are not close!
         >>> # The error message can also created at runtime by passing a callable.
         >>> def custom_msg(actual, expected, diagnostic_info):
+        ...     ratio = diagnostic_info.total_mismatches / diagnostic_info.number_of_elements
         ...     return (
         ...         f"Argh, we found {diagnostic_info.total_mismatches} mismatches! "
-        ...         f"That is {diagnostic_info.mismatch_ratio:.1%}!"
+        ...         f"That is {ratio:.1%}!"
         ...     )
         >>> torch.testing.assert_close(actual, expected, msg=custom_msg)
         AssertionError: Argh, we found 2 mismatches! That is 66.7%!
