Merge pull request #21666 from asmeurer/array_api-unique

seberg · web-flow · commit e5fcb9d52c1b · 2022-06-03T18:32:19.000-07:00
MAINT: Change array API unique_*() functions to not compare nans as equal
diff --git a/numpy/array_api/_set_functions.py b/numpy/array_api/_set_functions.py
@@ -46,6 +46,7 @@ def unique_all(x: Array, /) -> UniqueAllResult:
         return_counts=True,
         return_index=True,
         return_inverse=True,
+        equal_nan=False,
     )
     # np.unique() flattens inverse indices, but they need to share x's shape
     # See https://github.com/numpy/numpy/issues/20638
@@ -64,6 +65,7 @@ def unique_counts(x: Array, /) -> UniqueCountsResult:
         return_counts=True,
         return_index=False,
         return_inverse=False,
+        equal_nan=False,
     )
 
     return UniqueCountsResult(*[Array._new(i) for i in res])
@@ -80,6 +82,7 @@ def unique_inverse(x: Array, /) -> UniqueInverseResult:
         return_counts=False,
         return_index=False,
         return_inverse=True,
+        equal_nan=False,
     )
     # np.unique() flattens inverse indices, but they need to share x's shape
     # See https://github.com/numpy/numpy/issues/20638
@@ -98,5 +101,6 @@ def unique_values(x: Array, /) -> Array:
         return_counts=False,
         return_index=False,
         return_inverse=False,
+        equal_nan=False,
     )
     return Array._new(res)
