Repository: cupy/cupy

Issue:
[Bug, Enhancement] array_equal returns False on identical arrays with NaN values
* Conditions (you can just paste the output of `python -c 'import cupy; cupy.show_config()'`)
  - CuPy Version          : 8.0.0
  - OS/Platform : Pop OS
  - CUDA Version    : 10000
  - cuDNN Version         : 7605
  - NCCL Version    : 2604

* Code to reproduce
```python
import cupy as cp
a = cp.array([0, 1, cp.nan])
print(cp.array_equal(a, a))         # False
```
This relates exactly with https://github.com/numpy/numpy/issues/9229
This could be resolved by following the same procedure, resolved by https://github.com/numpy/numpy/pull/16128, which would essentially add an `equal_nan` kwarg to let user decide if NaN values should be considered equal or not, which is consistent with: https://github.com/cupy/cupy/blob/0dc6bbafb0cf75bb873531cbf69c6995a72c61a3/cupy/_logic/comparison.py#L60
https://github.com/cupy/cupy/blob/0dc6bbafb0cf75bb873531cbf69c6995a72c61a3/cupy/_logic/comparison.py#L89

* Error messages, stack traces, or logs
`<null>`


Discussion:

Resolving patch:
diff --git a/cupy/_logic/comparison.py b/cupy/_logic/comparison.py
--- a/cupy/_logic/comparison.py
+++ b/cupy/_logic/comparison.py
@@ -2,6 +2,7 @@
 
 import cupy
 from cupy import core
+from cupy._logic import content
 
 
 _is_close = core.create_ufunc(
@@ -38,12 +39,14 @@
 )
 
 
-def array_equal(a1, a2):
+def array_equal(a1, a2, equal_nan=False):
     """Returns ``True`` if two arrays are element-wise exactly equal.
 
     Args:
         a1 (cupy.ndarray): Input array to compare.
         a2 (cupy.ndarray): Input array to compare.
+        equal_nan (bool): If ``True``, NaN's in ``a1`` will be considered equal
+            to NaN's in ``a2``.
 
     Returns:
         cupy.ndarray: A boolean 0-dim array.
@@ -54,7 +57,15 @@ def array_equal(a1, a2):
     """
     if a1.shape != a2.shape:
         return cupy.array(False)
-    return (a1 == a2).all()
+    if not equal_nan:
+        return (a1 == a2).all()
+    # Handling NaN values if equal_nan is True
+    a1nan, a2nan = content.isnan(a1), content.isnan(a2)
+    # NaN's occur at different locations
+    if not (a1nan == a2nan).all():
+        return cupy.array(False)
+    # Shapes of a1, a2 and masks are guaranteed to be consistent by this point
+    return (a1[~a1nan] == a2[~a1nan]).all()
 
 
 def allclose(a, b, rtol=1.e-5, atol=1.e-8, equal_nan=False):


Test patch:
diff --git a/tests/cupy_tests/logic_tests/test_comparison.py b/tests/cupy_tests/logic_tests/test_comparison.py
--- a/tests/cupy_tests/logic_tests/test_comparison.py
+++ b/tests/cupy_tests/logic_tests/test_comparison.py
@@ -97,6 +97,27 @@ def test_array_equal_diff_length(self, xp, dtype):
         b = xp.array([1, 2, 3], dtype=dtype)
         return xp.array_equal(a, b)
 
+    @testing.with_requires('numpy>=1.19')
+    @testing.for_float_dtypes()
+    @testing.numpy_cupy_equal()
+    def test_array_equal_infinite_equal_nan(self, xp, dtype):
+        nan = float('nan')
+        inf = float('inf')
+        ninf = float('-inf')
+        a = xp.array([0, nan, inf, ninf], dtype=dtype)
+        b = xp.array([0, nan, inf, ninf], dtype=dtype)
+        return xp.array_equal(a, b, equal_nan=True)
+
+    @testing.with_requires('numpy>=1.19')
+    @testing.for_complex_dtypes()
+    @testing.numpy_cupy_equal()
+    def test_array_equal_complex_equal_nan(self, xp, dtype):
+        a = xp.array([1+2j], dtype=dtype)
+        b = a.copy()
+        b.imag = xp.nan
+        a.real = xp.nan
+        return xp.array_equal(a, b, equal_nan=True)
+
     @testing.numpy_cupy_equal()
     def test_array_equal_diff_dtypes_not_equal(self, xp):
         a = xp.array([0.9e-5, 1.1e-5, 100.5, 10.5])