up the priority of numpy array comparisons in self.assertEqual (#59067)

Summary:
Fixes #58988.

Pull Request resolved: #59067

Reviewed By: jbschlosser

Differential Revision: D28986642

Pulled By: heitorschueroff

fbshipit-source-id: 3ef2d26b4010fc3519d0a1a020ea446ffeb46ba0

test/test_binary_ufuncs.py

@@ -1262,10 +1262,10 @@ def test_maximum_minimum_float(self, device, dtype):
 
             if alias is not None:
                 alias_result = alias(a_tensor, b_tensor)
-                self.assertEqual(alias_result, tensor_result)
+                self.assertEqual(alias_result, tensor_result, exact_dtype=False)
 
-            self.assertEqual(tensor_result, numpy_result)
-            self.assertEqual(out, numpy_result)
+            self.assertEqual(tensor_result, numpy_result, exact_dtype=False)
+            self.assertEqual(out, numpy_result, exact_dtype=False)
 
     @dtypes(*(torch.testing.get_all_fp_dtypes()))
     def test_maximum_minimum_float_nan_and_inf(self, device, dtype):
@@ -1721,7 +1721,7 @@ def test_hypot(self, device, dtype):
                 expected = torch.sqrt(input[0] * input[0] + input[1] * input[1])
             else:
                 expected = np.hypot(input[0].cpu().numpy(), input[1].cpu().numpy())
-            self.assertEqual(actual, expected)
+            self.assertEqual(actual, expected, exact_dtype=False)
 
     @onlyOnCPUAndCUDA
     @dtypes(torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64)
@@ -1731,14 +1731,15 @@ def test_gcd(self, device, dtype):
         t2 = torch.tensor([0, 0, 10], dtype=dtype, device=device)
         actual = torch.gcd(t1, t2)
         expected = np.gcd([0, 10, 0], [0, 0, 10])
-        self.assertEqual(actual, expected)
+        self.assertEqual(actual, expected, exact_dtype=False)
 
         if dtype == torch.uint8:
             # Test unsigned integers with potential sign issues (i.e., uint8 with value >= 128)
             a = torch.tensor([190, 210], device=device, dtype=dtype)
             b = torch.tensor([190, 220], device=device, dtype=dtype)
             actual = torch.gcd(a, b)
             expected = torch.tensor([190, 10], device=device, dtype=dtype)
+            self.assertEqual(actual, expected)
         else:
             # Compares with NumPy
             a = torch.randint(-20, 20, (1024,), device=device, dtype=dtype)
@@ -1755,14 +1756,14 @@ def test_lcm(self, device, dtype):
         t2 = torch.tensor([0, 0, 10], dtype=dtype, device=device)
         actual = torch.lcm(t1, t2)
         expected = np.lcm([0, 10, 0], [0, 0, 10])
-        self.assertEqual(actual, expected)
+        self.assertEqual(actual, expected, exact_dtype=False)
 
         # Compares with NumPy
         a = torch.randint(-20, 20, (1024,), device=device, dtype=dtype)
         b = torch.randint(-20, 20, (1024,), device=device, dtype=dtype)
         actual = torch.lcm(a, b)
         expected = np.lcm(a.cpu().numpy(), b.cpu().numpy())
-        self.assertEqual(actual, expected)
+        self.assertEqual(actual, expected, exact_dtype=False)
 
     @onlyOnCPUAndCUDA
     @dtypes(torch.float32, torch.float64)
@@ -2490,14 +2491,14 @@ def test_dx(sizes, dim, dx, device):
             actual = torch.trapz(t, dx=dx, dim=dim)
             expected = np.trapz(t.cpu().numpy(), dx=dx, axis=dim)
             self.assertEqual(expected.shape, actual.shape)
-            self.assertEqual(expected, actual)
+            self.assertEqual(expected, actual, exact_dtype=False)
 
         def test_x(sizes, dim, x, device):
             t = torch.randn(sizes, device=device)
             actual = torch.trapz(t, x=torch.tensor(x, device=device), dim=dim)
             expected = np.trapz(t.cpu().numpy(), x=x, axis=dim)
             self.assertEqual(expected.shape, actual.shape)
-            self.assertEqual(expected, actual.cpu())
+            self.assertEqual(expected, actual.cpu(), exact_dtype=False)
 
         test_dx((2, 3, 4), 1, 1, device)
         test_dx((10, 2), 0, 0.1, device)

test/test_linalg.py

@@ -95,25 +95,27 @@ def run_test_case(a, b):
             if dtype == torch.bfloat16:
                 a_np = a.to(torch.double).cpu().numpy()
                 b_np = b.to(torch.double).cpu().numpy()
+                exact_dtype = False
             else:
                 a_np = a.cpu().numpy()
                 b_np = b.cpu().numpy()
+                exact_dtype = True
             expected = np.outer(a_np, b_np)
 
-            self.assertEqual(torch.outer(a, b), expected)
-            self.assertEqual(torch.Tensor.outer(a, b), expected)
+            self.assertEqual(torch.outer(a, b), expected, exact_dtype=False)
+            self.assertEqual(torch.Tensor.outer(a, b), expected, exact_dtype=False)
 
-            self.assertEqual(torch.ger(a, b), expected)
-            self.assertEqual(torch.Tensor.ger(a, b), expected)
+            self.assertEqual(torch.ger(a, b), expected, exact_dtype=False)
+            self.assertEqual(torch.Tensor.ger(a, b), expected, exact_dtype=False)
 
             # test out variant
             out = torch.empty(a.size(0), b.size(0), device=device, dtype=dtype)
             torch.outer(a, b, out=out)
-            self.assertEqual(out, expected)
+            self.assertEqual(out, expected, exact_dtype=False)
 
             out = torch.empty(a.size(0), b.size(0), device=device, dtype=dtype)
             torch.ger(a, b, out=out)
-            self.assertEqual(out, expected)
+            self.assertEqual(out, expected, exact_dtype=False)
 
         a = torch.randn(50).to(device=device, dtype=dtype)
         b = torch.randn(50).to(device=device, dtype=dtype)
@@ -184,7 +186,9 @@ def select_if_not_empty(t, i):
                     # see https://github.com/pytorch/pytorch/issues/56483
                     if m > n:
                         if torch.all(rank_1d == n):
-                            self.assertEqual(residuals, select_if_not_empty(residuals_2d, i), atol=1e-5, rtol=1e-5)
+                            self.assertEqual(
+                                residuals, select_if_not_empty(residuals_2d, i), atol=1e-5, rtol=1e-5, exact_dtype=False
+                            )
                         else:
                             self.assertTrue(residuals_2d.numel() == 0)
 
@@ -725,22 +729,24 @@ def check(m, a, b, beta, alpha):
                 a_np = a.to(torch.double).cpu().numpy()
                 b_np = b.to(torch.double).cpu().numpy()
                 m_np = m.to(torch.double).cpu().numpy()
+                exact_dtype = False
             else:
                 a_np = a.cpu().numpy()
                 b_np = b.cpu().numpy()
                 m_np = m.cpu().numpy()
+                exact_dtype = True
             if beta == 0:
                 expected = alpha * np.outer(a_np, b_np)
             else:
                 expected = beta * m_np + alpha * np.outer(a_np, b_np)
 
             res = torch.addr(m, a, b, beta=beta, alpha=alpha)
-            self.assertEqual(res, expected)
+            self.assertEqual(res, expected, exact_dtype=exact_dtype)
 
             # Test out variant
             out = torch.empty_like(res)
             torch.addr(m, a, b, beta=beta, alpha=alpha, out=out)
-            self.assertEqual(out, expected)
+            self.assertEqual(out, expected, exact_dtype=exact_dtype)
 
         m = make_tensor((50, 50), device=device, dtype=dtype, low=-2, high=2)
         a = make_tensor((50,), device=device, dtype=dtype, low=-2, high=2)
@@ -1591,7 +1597,7 @@ def test_cond(self, device, dtype):
         def run_test_case(input, p):
             result = torch.linalg.cond(input, p)
             result_numpy = np.linalg.cond(input.cpu().numpy(), p)
-            self.assertEqual(result, result_numpy, rtol=1e-2, atol=self.precision)
+            self.assertEqual(result, result_numpy, rtol=1e-2, atol=self.precision, exact_dtype=False)
             self.assertEqual(result.shape, result_numpy.shape)
 
             # test out= variant
@@ -1770,12 +1776,12 @@ def gen_error_message(input_size, ord, keepdim, dim=None):
                 expected = np.linalg.norm(xn, ord, keepdims=keepdim)
                 msg = gen_error_message(x.size(), ord, keepdim)
                 self.assertEqual(res.shape, expected.shape, msg=msg)
-                self.assertEqual(res, expected, msg=msg)
+                self.assertEqual(res, expected, msg=msg, exact_dtype=False)
 
                 res_out = torch.tensor([]).to(device)
                 torch.linalg.norm(x, ord, keepdim=keepdim, out=res_out)
                 self.assertEqual(res_out.shape, expected.shape, msg=msg)
-                self.assertEqual(res_out.cpu(), expected, msg=msg)
+                self.assertEqual(res_out.cpu(), expected, msg=msg, exact_dtype=False)
 
             # matrix norm
             x = torch.randn(25, 25, device=device, dtype=dtype)
@@ -1785,12 +1791,12 @@ def gen_error_message(input_size, ord, keepdim, dim=None):
                 expected = np.linalg.norm(xn, ord, keepdims=keepdim)
                 msg = gen_error_message(x.size(), ord, keepdim)
                 self.assertEqual(res.shape, expected.shape, msg=msg)
-                self.assertEqual(res, expected, msg=msg)
+                self.assertEqual(res, expected, msg=msg, exact_dtype=False)
 
                 res_out = torch.tensor([]).to(device)
                 torch.linalg.norm(x, ord, keepdim=keepdim, out=res_out)
                 self.assertEqual(res_out.shape, expected.shape, msg=msg)
-                self.assertEqual(res_out.cpu(), expected, msg=msg)
+                self.assertEqual(res_out.cpu(), expected, msg=msg, exact_dtype=False)
 
     # Test that linal.vector_norm gives the same result as numpy when inputs
     # contain extreme values (inf, -inf, nan)
@@ -1964,6 +1970,11 @@ def run_test_case(input, ord, dim, keepdim, should_error):
             for input_size, error_ords, dim in test_cases:
                 input = torch.randn(*input_size, dtype=dtype, device=device)
                 for ord in ord_matrix:
+                    # TODO: This was disabled in https://github.com/pytorch/pytorch/pull/59067, because it failed after
+                    #  fixing the underlying comparison mechanism. It should have never been passing in the first place.
+                    #  This should be reinstated as soon as https://github.com/pytorch/pytorch/issues/59198 is fixed.
+                    if input_size in ((S, S, 0), (1, S, 0)) and not keepdim and ord in (2, -2):
+                        continue
                     run_test_case(input, ord, dim, keepdim, ord in error_ords)
 
     def test_norm_fastpaths(self, device):
@@ -2169,8 +2180,8 @@ def run_test(shape, *, symmetric=False):
                 np.take_along_axis(actual_np[0], ind, axis=-1),
                 np.take_along_axis(actual_np[1], ind[:, None], axis=-1))
 
-            self.assertEqual(expected[0], sorted_actual[0])
-            self.assertEqual(abs(expected[1]), abs(sorted_actual[1]))
+            self.assertEqual(expected[0], sorted_actual[0], exact_dtype=False)
+            self.assertEqual(abs(expected[1]), abs(sorted_actual[1]), exact_dtype=False)
 
         shapes = [(0, 0),  # Empty matrix
                   (5, 5),  # Single matrix
@@ -2325,7 +2336,7 @@ def run_test(shape, *, symmetric=False):
             actual_np = actual.cpu().numpy()
             sorted_actual = np.take_along_axis(actual_np, ind, axis=-1)
 
-            self.assertEqual(expected, sorted_actual)
+            self.assertEqual(expected, sorted_actual, exact_dtype=False)
 
         shapes = [(0, 0),  # Empty matrix
                   (5, 5),  # Single matrix

test/test_nn.py

@@ -8269,7 +8269,7 @@ def _gelu_ref(X):
                     X = torch.rand(n, m, dtype=dtype, requires_grad=True, device=d)[:, ::2]
                 res = F.gelu(X)
                 ref = _gelu_ref(X.to(numpy_dtype).cpu().detach().numpy())
-                self.assertEqual(res, ref, rtol=rtol, atol=atol)
+                self.assertEqual(res, ref, rtol=rtol, atol=atol, exact_dtype=False)
                 if dtype == torch.float64:
                     gradcheck(F.gelu, [X], eps=1e-4)
 
@@ -12135,7 +12135,7 @@ def test_affine_2d_rotateRandom(self, device):
             for r in range(affine_tensor.size(1)):
                 for c in range(affine_tensor.size(2)):
                     grid_out = np.dot(grid_ary, [r, c, 1])
-                    self.assertEqual(affine_tensor[0, r, c], grid_out[:2])
+                    self.assertEqual(affine_tensor[0, r, c], grid_out[:2], exact_dtype=False)
 
             self.assertEqual(scipy_ary, gridsample_ary.reshape_as(scipy_ary))
 
@@ -12190,7 +12190,7 @@ def test_affine_3d_rotateRandom(self, device):
                 for r in range(affine_tensor.size(2)):
                     for c in range(affine_tensor.size(3)):
                         grid_out = np.dot(grid_ary, [i, r, c, 1])
-                        self.assertEqual(affine_tensor[0, i, r, c], grid_out[:3])
+                        self.assertEqual(affine_tensor[0, i, r, c], grid_out[:3], exact_dtype=False)
 
             self.assertEqual(scipy_ary, gridsample_ary.reshape_as(scipy_ary))
 

test/test_reductions.py

@@ -1814,7 +1814,7 @@ def test_median_real_values(self, device, dtype):
                 if size % 2 == 1:
                     # We can only test agains numpy for odd reductions because numpy
                     # returns the mean of the two medians and torch returns the lower
-                    self.assertEqual(res[0].cpu().numpy(), np.median(t_numpy, dim, keepdims=True))
+                    self.assertEqual(res[0].cpu().numpy(), np.median(t_numpy, dim, keepdims=True), exact_dtype=False)
 
     @dtypes(torch.float, torch.double)
     @dtypesIfCUDA(torch.half, torch.float, torch.double)
@@ -2086,7 +2086,7 @@ def _compare_std_var_with_numpy(self, op, device, dtype, input, dim,
                 return
             self.fail("Failed to hit RuntimeError!")
 
-        exact_dtype = input.dtype != torch.bfloat16
+        exact_dtype = input.dtype not in (torch.bfloat16, torch.complex32, torch.complex64, torch.complex128)
         self.assertEqual(torch_result, numpy_result, exact_dtype=exact_dtype)
 
     @dtypes(torch.float, torch.double, torch.cfloat, torch.cdouble)
@@ -2562,21 +2562,21 @@ def test_tensor_compare_ops_empty(self, device):
             error_msg = f"test function: {name}"
             self.assertEqual(torch.empty((2, 0), device=device), fn(master_input, dim=2), msg=error_msg)
             self.assertEqual(np_function(np_input, axis=2),
-                             fn(master_input, dim=2).cpu().numpy(), msg=error_msg)
+                             fn(master_input, dim=2).cpu().numpy(), msg=error_msg, exact_dtype=False)
 
             self.assertEqual(torch.empty((2, 0), device=device), fn(master_input, dim=-1), msg=error_msg)
             self.assertEqual(np_function(np_input, axis=-1),
-                             fn(master_input, dim=-1).cpu().numpy(), msg=error_msg)
+                             fn(master_input, dim=-1).cpu().numpy(), msg=error_msg, exact_dtype=False)
 
             self.assertEqual(torch.empty((2, 0, 1), device=device), fn(master_input, dim=2, keepdim=True),
                              msg=error_msg)
             self.assertEqual(np_function(np_input, axis=2, keepdims=True),
-                             fn(master_input, dim=2, keepdim=True).cpu().numpy(), msg=error_msg)
+                             fn(master_input, dim=2, keepdim=True).cpu().numpy(), msg=error_msg, exact_dtype=False)
 
             self.assertEqual(torch.empty((2, 0, 1), device=device), fn(master_input, dim=-1, keepdim=True),
                              msg=error_msg)
             self.assertEqual(np_function(np_input, axis=-1, keepdims=True),
-                             fn(master_input, dim=-1, keepdim=True).cpu().numpy(), msg=error_msg)
+                             fn(master_input, dim=-1, keepdim=True).cpu().numpy(), msg=error_msg, exact_dtype=False)
 
             # Check if function raises error on specified zero'd dimension as reduction dim.
             self.assertRaisesRegex(IndexError, "Expected reduction dim", lambda: fn(master_input, dim=1))
@@ -2643,18 +2643,20 @@ def test_tensor_reduce_ops_empty(self, device):
             # Check if reduction happens along the specified dimension.
             error_msg = f"test function: {name}"
             self.assertEqual(torch.empty((2, 0), device=device), fn(master_input, dim=2), msg=error_msg)
-            self.assertEqual(np_function(np_input, axis=2), fn(master_input, dim=2).cpu().numpy(), msg=error_msg)
+            self.assertEqual(np_function(np_input, axis=2), fn(master_input, dim=2).cpu().numpy(), msg=error_msg,
+                             exact_dtype=False)
 
             self.assertEqual(torch.empty((2, 0), device=device), fn(master_input, dim=-1), msg=error_msg)
-            self.assertEqual(np_function(np_input, axis=-1), fn(master_input, dim=-1).cpu().numpy(), msg=error_msg)
+            self.assertEqual(np_function(np_input, axis=-1), fn(master_input, dim=-1).cpu().numpy(), msg=error_msg,
+                             exact_dtype=False)
 
             self.assertEqual(torch.empty((2, 0, 1), device=device), fn(master_input, dim=2, keepdim=True), msg=error_msg)
             self.assertEqual(np_function(np_input, axis=2, keepdims=True), fn(master_input, dim=2, keepdim=True),
-                             msg=error_msg)
+                             msg=error_msg, exact_dtype=False)
 
             self.assertEqual(torch.empty((2, 0, 1), device=device), fn(master_input, dim=-1, keepdim=True), msg=error_msg)
             self.assertEqual(np_function(np_input, axis=-1, keepdims=True), fn(master_input, dim=-1, keepdim=True),
-                             msg=error_msg)
+                             msg=error_msg, exact_dtype=False)
 
             # Check if returned data is correct.
             check_func = (torch.testing.assert_allclose if math.isnan(return_value) or math.isinf(return_value) else