Call a separate function to ensure same rank of arrays

sklam · Apr 21, 2022 · Mar 26, 2022 · Mar 26, 2022 · Apr 1, 2022 · Apr 2, 2022
commit f09d8ffd2555e59a70ec0618b1532e8d95511ea2
diff --git a/numba/np/arraymath.py b/numba/np/arraymath.py
@@ -921,6 +921,17 @@ def allclose_scalars(a_v, b_v, rtol=1e-05, atol=1e-08, equal_nan=False):
     return True
 
 
+@register_jitable
+def np_allclose_impl_array_array_no_broadcast(a, b, rtol=1e-05, atol=1e-08,
+                                              equal_nan=False):
+    for av, bv in np.nditer((a, b)):
+        if not allclose_scalars(av.item(), bv.item(), rtol=rtol,
+                                atol=atol, equal_nan=equal_nan):
+            return False
+
+    return True
+
+
 @overload(np.allclose)
 @overload_method(types.Array, "allclose")
 def np_allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False):
@@ -979,19 +990,14 @@ def np_allclose_impl_array_array(a, b, rtol=1e-05, atol=1e-08,
             b = np.asarray(b)
             if a.shape != b.shape:
                 if _can_broadcast_returns_bool(a, b.shape):
-                    a = np.broadcast_to(a, b.shape)
+                    return np_allclose_impl_array_array_no_broadcast(np.broadcast_to(a, b.shape), b, rtol=rtol, atol=atol, equal_nan=equal_nan)
                 elif _can_broadcast_returns_bool(b, a.shape):
-                    b = np.broadcast_to(b, a.shape)
+                    return np_allclose_impl_array_array_no_broadcast(a, np.broadcast_to(b, a.shape), rtol=rtol, atol=atol, equal_nan=equal_nan)
                 else:
                     raise ValueError('operands could not be broadcast together '
                                      'with remapped shapes')
 
-            for av, bv in zip(np.nditer(a), np.nditer(b)):
-                if not allclose_scalars(av.item(), bv.item(), rtol=rtol,
-                                        atol=atol, equal_nan=equal_nan):
-                    return False
-
-            return True
+            return np_allclose_impl_array_array_no_broadcast(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)
 
         return np_allclose_impl_array_array
 

diff --git a/numba/tests/test_np_functions.py b/numba/tests/test_np_functions.py
@@ -3250,7 +3250,7 @@ def __new__(cls, *args, **kwargs):
         self.assertTrue(type(np.allclose(a, a)) is bool)
 
         noise_levels = [1.0, 1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6, 0.0]
-        zero_array = np.zeros((25,))
+        zero_array = np.zeros((25, 4))
         a = np.random.ranf((25, 4))
         for noise in noise_levels:
             for rtol in noise_levels: