Add dpnp.top_k

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 * Added support for boolean input in `dpnp.positive` and the `+` operator of `dpnp.ndarray` [#3128](https://github.com/IntelPython/dpnp/pull/3128)
+* Added `dpnp.top_k` to return the `k` largest or smallest elements and their indices along an axis, with a `sorted` keyword to return them in sorted order [#3133](https://github.com/IntelPython/dpnp/pull/3133)
 
 ### Changed
 
@@ -23,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Deprecated `dpnp.dot` and `dpnp.ndarray.dot` when `a.ndim >= 2` and `b.ndim > 2`, which place the leading dimensions of `b` in the middle of the result; use `dpnp.tensordot(a, b, axes=(-1, -2))` or `dpnp.matmul` instead [#3124](https://github.com/IntelPython/dpnp/pull/3124)
 * Deprecated the upcast of `arr` by a Python numeric scalar in `dpnp.append` with `axis=None` (e.g. `int8` array and `1` giving `int64`) [#3125](https://github.com/IntelPython/dpnp/pull/3125)
 * Deprecated `dpnp.isfortran`; use `a.flags.fnc` instead [#3140](https://github.com/IntelPython/dpnp/pull/3140)
+* Deprecated omitting `axis` in `dpnp.tensor.top_k` for an input that is not 1-D, which searches the flattened array; the default will change to `axis=-1` as in the Array API standard, so pass `axis=None` to keep the current behavior [#3133](https://github.com/IntelPython/dpnp/pull/3133)
 
 ### Removed
 

doc/reference/sort.rst

@@ -18,6 +18,7 @@ Sorting
    sort_complex
    partition
    argpartition
+   top_k
 
 .. seealso::
    :func:`dpnp.ndarray.sort`

dpnp/__init__.py

@@ -515,6 +515,7 @@
     partition,
     sort,
     sort_complex,
+    top_k,
 )
 
 # -----------------------------------------------------------------------------
@@ -995,6 +996,7 @@
     "searchsorted",
     "sort",
     "sort_complex",
+    "top_k",
     "where",
 ]
 

dpnp/dpnp_iface_sorting.py

@@ -39,7 +39,9 @@
 
 """
 
+import operator
 from collections.abc import Sequence
+from typing import NamedTuple
 
 import dpnp
 import dpnp.tensor as dpt
@@ -55,6 +57,12 @@
 from .tensor._numpy_helper import normalize_axis_index
 
 
+# pylint:disable=missing-class-docstring
+class TopKResult(NamedTuple):
+    values: dpnp.ndarray
+    indices: dpnp.ndarray
+
+
 def _wrap_sort_argsort(
     a,
     _sorting_fn,
@@ -448,3 +456,136 @@ def sort_complex(a):
             b_dt = map_dtype_to_device(dpnp.complex128, b.sycl_device)
         return b.astype(b_dt)
     return b
+
+
+def top_k(a, k, /, *, axis=-1, mode="largest", sorted=True):
+    """
+    Returns the `k` largest or smallest elements and their indices along an
+    axis.
+
+    For full documentation refer to :obj:`numpy.top_k`.
+
+    Parameters
+    ----------
+    a : {dpnp.ndarray, usm_ndarray}
+        Input array.
+    k : int
+        The number of largest or smallest elements to return. `k` must be
+        a non-negative integer not greater than the size of `a` along `axis`.
+    axis : {None, int}, optional
+        Axis along which to find the largest or smallest elements. If
+        ``None``, the flattened array is used.
+
+        Default: ``-1``.
+    mode : {"largest", "smallest"}, optional
+        If ``"largest"``, the largest elements are returned. If
+        ``"smallest"``, the smallest elements are returned. NaN values (and
+        complex values with a NaN component) are ordered to the end regardless
+        of `mode`, so they are only returned if they are among the top `k`
+        values.
+
+        Default: ``"largest"``.
+    sorted : bool, optional
+        If ``True``, the top `k` elements are returned in sorted order:
+        largest-first for ``mode="largest"`` and smallest-first for
+        ``mode="smallest"``. If ``False``, the order of the elements is
+        unspecified.
+
+        Default: ``True``.
+
+    Returns
+    -------
+    values : dpnp.ndarray
+        The top `k` values of `a`, with the same data type as `a`. The shape
+        is that of `a` with the dimension along `axis` replaced by `k`.
+    indices : dpnp.ndarray
+        The indices of `values` along `axis`, with the same shape as `values`
+        and the default array index data type.
+
+    See Also
+    --------
+    :obj:`dpnp.tensor.top_k` : Array API version; returns the elements in an
+        unspecified order.
+    :obj:`dpnp.partition` : Partial sort.
+    :obj:`dpnp.sort` : Full sorting.
+
+    Notes
+    -----
+    The returned indices are not guaranteed to be stable, i.e., the order of
+    the indices of equal values is not guaranteed to match their order in
+    `a`. This is the case regardless of the value of `sorted`.
+
+    In contrast to :obj:`numpy.top_k`, which returns empty arrays when `axis`
+    is empty, :obj:`dpnp.top_k` raises ``ValueError`` for an empty `axis` and
+    ``k > 0``, as it does for any `k` greater than the size of `axis`.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> a = np.array([[1, 2, 3, 4], [4, 3, 2, 1]])
+    >>> res = np.top_k(a, 2)
+    >>> res.values
+    array([[4, 3],
+           [4, 3]])
+    >>> res.indices
+    array([[3, 2],
+           [0, 1]])
+
+    >>> values, indices = np.top_k(a, 2, axis=0)
+    >>> values
+    array([[4, 3, 3, 4],
+           [1, 2, 2, 1]])
+    >>> indices
+    array([[1, 1, 0, 0],
+           [0, 0, 1, 1]])
+
+    >>> values, indices = np.top_k(a, 2, axis=1, mode="smallest")
+    >>> values
+    array([[1, 2],
+           [1, 2]])
+    >>> indices
+    array([[0, 1],
+           [3, 2]])
+
+    >>> values, indices = np.top_k(np.array([1.0, 2.0, 3.0, np.nan]), 2)
+    >>> values
+    array([3., 2.])
+    >>> indices
+    array([2, 1])
+
+    """
+
+    if mode not in ("largest", "smallest"):
+        raise ValueError(
+            f'mode(="{mode}") must be either "largest" or "smallest".'
+        )
+
+    usm_a = dpnp.get_usm_ndarray(a)
+    if axis is None:
+        usm_a = dpt.reshape(usm_a, (-1,))
+        axis = 0
+    else:
+        axis = normalize_axis_index(axis, ndim=usm_a.ndim)
+
+    k = operator.index(k)
+    n = usm_a.shape[axis]
+    descending = mode == "largest"
+
+    if sorted and 1 < k <= n < 2 * k:
+        # for large k a full sort is cheaper than selecting and then sorting
+        usm_inds = dpt.argsort(usm_a, axis=axis, descending=descending)
+        usm_inds = usm_inds[(slice(None),) * axis + (slice(k),)]
+        usm_vals = dpt.take_along_axis(usm_a, usm_inds, axis=axis)
+    else:
+        usm_vals, usm_inds = dpt.top_k(usm_a, k, axis=axis, mode=mode)
+
+        # tensor.top_k returns the elements in an unspecified order
+        if sorted and k > 1:
+            sort_inds = dpt.argsort(usm_vals, axis=axis, descending=descending)
+            usm_vals = dpt.take_along_axis(usm_vals, sort_inds, axis=axis)
+            usm_inds = dpt.take_along_axis(usm_inds, sort_inds, axis=axis)
+
+    return TopKResult(
+        dpnp_array._create_from_usm_ndarray(usm_vals),
+        dpnp_array._create_from_usm_ndarray(usm_inds),
+    )

dpnp/tensor/_sorting.py

@@ -27,6 +27,7 @@
 # *****************************************************************************
 
 import operator
+import warnings
 from typing import NamedTuple
 
 import dpctl.utils as du
@@ -299,7 +300,11 @@ class TopKResult(NamedTuple):
     indices: dpt.usm_ndarray
 
 
-def top_k(x, k, /, *, axis=None, mode="largest"):
+# TODO: default to `axis=-1` once the deprecation period ends
+_NO_AXIS = object()
+
+
+def top_k(x, k, /, *, axis=_NO_AXIS, mode="largest"):
     """top_k(x, k, axis=None, mode="largest")
 
     Returns the `k` largest or smallest values and their indices in the input
@@ -309,10 +314,16 @@ def top_k(x, k, /, *, axis=None, mode="largest"):
         x (usm_ndarray):
             input array.
         k (int):
-            number of elements to find. Must be a positive integer value.
+            number of elements to find. Must be a non-negative integer value.
         axis (Optional[int]):
             axis along which to search. If `None`, the search will be performed
             over the flattened array. Default: ``None``.
+
+            .. deprecated:: 0.22.0
+                Omitting `axis` for an input that is not 1-D is deprecated.
+                The default will change to ``-1``, as in the Array API
+                standard. Pass ``axis=None`` to keep searching the flattened
+                array.
         mode (Literal["largest", "smallest"]):
             search mode. Must be one of the following modes:
 
@@ -350,9 +361,21 @@ def top_k(x, k, /, *, axis=None, mode="largest"):
 
     k = operator.index(k)
     if k < 0:
-        raise ValueError("`k` must be a positive integer value")
+        raise ValueError("`k` must be a non-negative integer value")
 
     nd = x.ndim
+    if axis is _NO_AXIS:
+        if nd != 1:
+            warnings.warn(
+                "Omitting `axis` in `dpnp.tensor.top_k` for an input that is "
+                "not 1-D is deprecated. The default will change from `None` "
+                "(search the flattened array) to `-1` (search the last axis), "
+                "as in the Array API standard. Pass `axis=None` to keep the "
+                "current behavior or `axis=-1` to adopt the new one.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        axis = None
     if axis is None:
         sz = x.size
         if nd == 0:

dpnp/tests/tensor/test_usm_ndarray_top_k.py

@@ -26,6 +26,9 @@
 # THE POSSIBILITY OF SUCH DAMAGE.
 # *****************************************************************************
 
+import math
+import warnings
+
 import numpy as np
 import pytest
 from numpy.testing import assert_array_equal
@@ -439,7 +442,7 @@ def test_top_k_0d():
     assert a.ndim == 0
     assert a.size == 1
 
-    r = dpt.top_k(a, 1)
+    r = dpt.top_k(a, 1, axis=None)
     assert r.values == a
     assert r.indices == dpt.zeros_like(a, dtype=r.indices.dtype)
 
@@ -477,8 +480,7 @@ def test_top_k_axis0():
 def test_top_k_validation():
     get_queue_or_skip()
     x = dpt.ones(10, dtype=dpt.int64)
-    with pytest.raises(ValueError):
-        # k must be positive
+    with pytest.raises(ValueError, match="non-negative"):
         dpt.top_k(x, -1)
     with pytest.raises(TypeError):
         # argument should be usm_ndarray
@@ -490,12 +492,38 @@ def test_top_k_validation():
         dpt.top_k(x2, 100, axis=1)
     with pytest.raises(ValueError):
         # for 0d arrays, k must be 1
-        dpt.top_k(x[0], 2)
+        dpt.top_k(x[0], 2, axis=None)
     with pytest.raises(ValueError):
         # mode must be "largest", or "smallest"
         dpt.top_k(x, 2, mode="invalid")
 
 
+@pytest.mark.parametrize("shape", [(), (4, 5), (2, 3, 4)])
+def test_top_k_default_axis_deprecated(shape):
+    get_queue_or_skip()
+    x = dpt.reshape(dpt.arange(math.prod(shape), dtype="i4"), shape)
+
+    with pytest.warns(DeprecationWarning, match="axis=None"):
+        r = dpt.top_k(x, 1)
+
+    # the input is still searched over the flattened array
+    expected = dpt.top_k(x, 1, axis=None)
+    assert dpt.all(r.values == expected.values)
+    assert dpt.all(r.indices == expected.indices)
+
+
+def test_top_k_axis_no_warning():
+    get_queue_or_skip()
+    x = dpt.reshape(dpt.arange(20, dtype="i4"), (4, 5))
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", DeprecationWarning)
+        dpt.top_k(x, 2, axis=None)
+        dpt.top_k(x, 2, axis=-1)
+        # the new default gives the same result for 1-D input
+        dpt.top_k(x[0], 2)
+
+
 @pytest.mark.parametrize("dtype", ["i4", "u4", "i8", "f4", "f8"])
 @pytest.mark.parametrize("mode", ["largest", "smallest"])
 def test_top_k_long_rows(dtype, mode):