numpy · ngoldbaum · Dec 5, 2023 · Nov 7, 2023
diff --git a/doc/release/upcoming_changes/25088.new_feature.rst b/doc/release/upcoming_changes/25088.new_feature.rst
@@ -0,0 +1,6 @@
+Array API set functions
+-----------------------
+
+`numpy.unique_all`, `numpy.unique_counts`, `numpy.unique_inverse`,
+and `numpy.unique_values` functions have been added for Array API compatiblity.
+They provide functionality of `numpy.unique` with different sets of flags.
diff --git a/doc/source/reference/array_api.rst b/doc/source/reference/array_api.rst
@@ -119,10 +119,6 @@ The following functions are named differently in the array API
    * - ``pow``
      - ``power``
      -
-   * - ``unique_all``, ``unique_counts``, ``unique_inverse``, and
-       ``unique_values``
-     - ``unique``
-     - Each is equivalent to ``np.unique`` with certain flags set.
 
 
 Function instead of method

diff --git a/doc/source/reference/routines.set.rst b/doc/source/reference/routines.set.rst
@@ -9,6 +9,10 @@ Making proper sets
    :toctree: generated/
 
    unique
+   unique_all
+   unique_counts
+   unique_inverse
+   unique_values
 
 Boolean operations
 ------------------

diff --git a/numpy/__init__.py b/numpy/__init__.py
@@ -208,7 +208,8 @@
         real_if_close, typename, mintypecode, common_type
     )
     from .lib._arraysetops_impl import (
-        ediff1d, in1d, intersect1d, isin, setdiff1d, setxor1d, union1d, unique
+        ediff1d, in1d, intersect1d, isin, setdiff1d, setxor1d, union1d,
+        unique, unique_all, unique_counts, unique_inverse, unique_values
     )
     from .lib._ufunclike_impl import fix, isneginf, isposinf
     from .lib._arraypad_impl import pad

diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
@@ -406,13 +406,17 @@ from numpy.lib._arraypad_impl import (
 
 from numpy.lib._arraysetops_impl import (
     ediff1d as ediff1d,
+    in1d as in1d,
     intersect1d as intersect1d,
+    isin as isin,
+    setdiff1d as setdiff1d,
     setxor1d as setxor1d,
     union1d as union1d,
-    setdiff1d as setdiff1d,
     unique as unique,
-    in1d as in1d,
-    isin as isin,
+    unique_all as unique_all,
+    unique_counts as unique_counts,
+    unique_inverse as unique_inverse,
+    unique_values as unique_values,
 )
 
 from numpy.lib._function_base_impl import (

diff --git a/numpy/lib/_arraysetops_impl.py b/numpy/lib/_arraysetops_impl.py
@@ -16,6 +16,7 @@
 """
 import functools
 import warnings
+from typing import NamedTuple
 
 import numpy as np
 from numpy._core import overrides
@@ -26,9 +27,10 @@
 
 
 __all__ = [
-    'ediff1d', 'intersect1d', 'setxor1d', 'union1d', 'setdiff1d', 'unique',
-    'in1d', 'isin'
-    ]
+    "ediff1d", "in1d", "intersect1d", "isin", "setdiff1d", "setxor1d",
+    "union1d", "unique", "unique_all", "unique_counts", "unique_inverse",
+    "unique_values"
+]
 
 
 def _ediff1d_dispatcher(ary, to_end=None, to_begin=None):
@@ -364,6 +366,188 @@ def _unique1d(ar, return_index=False, return_inverse=False,
     return ret
 
 
+# Array API set functions
+
+class UniqueAllResult(NamedTuple):
+    values: np.ndarray
+    indices: np.ndarray
+    inverse_indices: np.ndarray
+    counts: np.ndarray
+
+
+class UniqueCountsResult(NamedTuple):
+    values: np.ndarray
+    counts: np.ndarray
+
+
+class UniqueInverseResult(NamedTuple):
+    values: np.ndarray
+    inverse_indices: np.ndarray
+
+
+def _unique_all_dispatcher(x, /):
+    return (x,)
+
+
+@array_function_dispatch(_unique_all_dispatcher)
+def unique_all(x):
+    """
+    Returns the unique elements of an input array `x`, the first occurring
+    indices for each unique element in `x`, the indices from the set of unique
+    elements that reconstruct `x`, and the corresponding counts for each
+    unique element in `x`.
+
+    This function is Array API compatible alternative to:
+    ``np.unique(x, return_index=True, return_inverse=True, return_counts=True,
+    equal_nan=False)``
+
+    Parameters
+    ----------
+    x : array_like
+        Input array. It will be flattened if it is not already 1-D.
+
+    Returns
+    -------
+    out : namedtuple
+        The result containing:
+        * values - The unique elements of an input array.
+        * indices - The first occurring indices for each unique element.
+        * inverse_indices - The indices from the set of unique elements
+          that reconstruct `x`.
+        * counts - The corresponding counts for each unique element.
+
+    See Also
+    --------
+    unique : Find the unique elements of an array.
+
+    """
+    result = unique(
+        x,
+        return_index=True,
+        return_inverse=True,
+        return_counts=True,
+        equal_nan=False
+    )
+    return UniqueAllResult(*result)
+
+
+def _unique_counts_dispatcher(x, /):
+    return (x,)
+
+
+@array_function_dispatch(_unique_counts_dispatcher)
+def unique_counts(x):
+    """
+    Returns the unique elements of an input array `x`, and the corresponding
+    counts for each unique element in `x`.
+
+    This function is Array API compatible alternative to:
+    ``np.unique(x, return_counts=True, equal_nan=False)``
+
+    Parameters
+    ----------
+    x : array_like
+        Input array. It will be flattened if it is not already 1-D.
+
+    Returns
+    -------
+    out : namedtuple
+        The result containing:
+        * values - The unique elements of an input array.
+        * counts - The corresponding counts for each unique element.
+
+    See Also
+    --------
+    unique : Find the unique elements of an array.
+
+    """
+    result = unique(
+        x,
+        return_index=False,
+        return_inverse=False,
+        return_counts=True,
+        equal_nan=False
+    )
+    return UniqueCountsResult(*result)
+
+
+def _unique_inverse_dispatcher(x, /):
+    return (x,)
+
+
+@array_function_dispatch(_unique_inverse_dispatcher)
+def unique_inverse(x):
+    """
+    Returns the unique elements of an input array `x` and the indices
+    from the set of unique elements that reconstruct `x`.
+
+    This function is Array API compatible alternative to:
+    ``np.unique(x, return_inverse=True, equal_nan=False)``
+
+    Parameters
+    ----------
+    x : array_like
+        Input array. It will be flattened if it is not already 1-D.
+
+    Returns
+    -------
+    out : namedtuple
+        The result containing:
+        * values - The unique elements of an input array.
+        * inverse_indices - The indices from the set of unique elements
+          that reconstruct `x`.
+
+    See Also
+    --------
+    unique : Find the unique elements of an array.
+
+    """
+    result = unique(
+        x,
+        return_index=False,
+        return_inverse=True,
+        return_counts=False,
+        equal_nan=False
+    )
+    return UniqueInverseResult(*result)
+
+
+def _unique_values_dispatcher(x, /):
+    return (x,)
+
+
+@array_function_dispatch(_unique_values_dispatcher)
+def unique_values(x):
+    """
+    Returns the unique elements of an input array `x`.
+
+    This function is Array API compatible alternative to:
+    ``np.unique(x, equal_nan=False)``
+
+    Parameters
+    ----------
+    x : array_like
+        Input array. It will be flattened if it is not already 1-D.
+
+    Returns
+    -------
+    out : ndarray
+        The unique elements of an input array.
+
+    See Also
+    --------
+    unique : Find the unique elements of an array.
+
+    """
+    return unique(
+        x,
+        return_index=False,
+        return_inverse=False,
+        return_counts=False,
+        equal_nan=False
+    )
+
+
 def _intersect1d_dispatcher(
         ar1, ar2, assume_unique=None, return_indices=None):
     return (ar1, ar2)

diff --git a/numpy/lib/_arraysetops_impl.pyi b/numpy/lib/_arraysetops_impl.pyi
@@ -1,9 +1,11 @@
 from typing import (
-    Literal as L,
     Any,
-    TypeVar,
+    Generic,
+    Literal as L,
+    NamedTuple,
     overload,
     SupportsIndex,
+    TypeVar,
 )
 
 import numpy as np
@@ -85,6 +87,20 @@ _SCTNoCast = TypeVar(
     void,
 )
 
+class UniqueAllResult(NamedTuple, Generic[_SCT]):
+    values: NDArray[_SCT]
+    indices: NDArray[intp]
+    inverse_indices: NDArray[intp]
+    counts: NDArray[intp]
+
+class UniqueCountsResult(NamedTuple, Generic[_SCT]):
+    values: NDArray[_SCT]
+    counts: NDArray[intp]
+
+class UniqueInverseResult(NamedTuple, Generic[_SCT]):
+    values: NDArray[_SCT]
+    inverse_indices: NDArray[intp]
+
 __all__: list[str]
 
 @overload
@@ -279,6 +295,34 @@ def unique(
     equal_nan: bool = ...,
 ) -> tuple[NDArray[Any], NDArray[intp], NDArray[intp], NDArray[intp]]: ...
 
+@overload
+def unique_all(
+    x: _ArrayLike[_SCT], /
+) -> UniqueAllResult[_SCT]: ...
+@overload
+def unique_all(
+    x: ArrayLike, /
+) -> UniqueAllResult[Any]: ...
+
+@overload
+def unique_counts(
+    x: _ArrayLike[_SCT], /
+) -> UniqueCountsResult[_SCT]: ...
+@overload
+def unique_counts(
+    x: ArrayLike, /
+) -> UniqueCountsResult[Any]: ...
+
+@overload
+def unique_inverse(x: _ArrayLike[_SCT], /) -> UniqueInverseResult[_SCT]: ...
+@overload
+def unique_inverse(x: ArrayLike, /) -> UniqueInverseResult[Any]: ...
+
+@overload
+def unique_values(x: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ...
+@overload
+def unique_values(x: ArrayLike, /) -> NDArray[Any]: ...
+
 @overload
 def intersect1d(
     ar1: _ArrayLike[_SCTNoCast],

diff --git a/numpy/lib/tests/test_arraysetops.py b/numpy/lib/tests/test_arraysetops.py
@@ -919,3 +919,34 @@ def test_unique_nanequals(self):
         not_unq = np.unique(a, equal_nan=False)
         assert_array_equal(unq, np.array([1, np.nan]))
         assert_array_equal(not_unq, np.array([1, np.nan, np.nan, np.nan]))
+
+    def test_unique_array_api_functions(self):
+        arr = np.array([np.nan, 1, 4, 1, 3, 4, np.nan, 5, 1])
+
+        for res_unique_array_api, res_unique in [
+            (
+                np.unique_values(arr),
+                np.unique(arr, equal_nan=False)
+            ),
+            (
+                np.unique_counts(arr),
+                np.unique(arr, return_counts=True, equal_nan=False)
+            ),
+            (
+                np.unique_inverse(arr),
+                np.unique(arr, return_inverse=True, equal_nan=False)
+            ),
+            (
+                np.unique_all(arr),
+                np.unique(
+                    arr,
+                    return_index=True,
+                    return_inverse=True,
+                    return_counts=True,
+                    equal_nan=False
+                )
+            )
+        ]:
+            assert len(res_unique_array_api) == len(res_unique)
+            for actual, expected in zip(res_unique_array_api, res_unique):
+                assert_array_equal(actual, expected)
diff --git a/numpy/typing/tests/data/reveal/arraysetops.pyi b/numpy/typing/tests/data/reveal/arraysetops.pyi
@@ -3,6 +3,9 @@ from typing import Any
 
 import numpy as np
 import numpy.typing as npt
+from numpy.lib._arraysetops_impl import (
+    UniqueAllResult, UniqueCountsResult, UniqueInverseResult
+)
 
 if sys.version_info >= (3, 11):
     from typing import assert_type
@@ -61,3 +64,12 @@ assert_type(np.unique(AR_f8, return_inverse=True, return_counts=True), tuple[npt
 assert_type(np.unique(AR_LIKE_f8, return_inverse=True, return_counts=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp]])
 assert_type(np.unique(AR_f8, return_index=True, return_inverse=True, return_counts=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp], npt.NDArray[np.intp]])
 assert_type(np.unique(AR_LIKE_f8, return_index=True, return_inverse=True, return_counts=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp], npt.NDArray[np.intp]])
+
+assert_type(np.unique_all(AR_f8), UniqueAllResult[np.float64])
+assert_type(np.unique_all(AR_LIKE_f8), UniqueAllResult[Any])
+assert_type(np.unique_counts(AR_f8), UniqueCountsResult[np.float64])
+assert_type(np.unique_counts(AR_LIKE_f8), UniqueCountsResult[Any])
+assert_type(np.unique_inverse(AR_f8), UniqueInverseResult[np.float64])
+assert_type(np.unique_inverse(AR_LIKE_f8), UniqueInverseResult[Any])
+assert_type(np.unique_values(AR_f8), npt.NDArray[np.float64])
+assert_type(np.unique_values(AR_LIKE_f8), npt.NDArray[Any])