diff --git a/CHANGELOG.md b/CHANGELOG.md index 33e1bd4aa64..8acc4098d6a 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/doc/reference/sort.rst b/doc/reference/sort.rst index 2f73aa77540..0711f8e5485 100644 --- a/doc/reference/sort.rst +++ b/doc/reference/sort.rst @@ -18,6 +18,7 @@ Sorting sort_complex partition argpartition + top_k .. seealso:: :func:`dpnp.ndarray.sort` diff --git a/dpnp/__init__.py b/dpnp/__init__.py index 68a683d8095..ce3b196e070 100644 --- a/dpnp/__init__.py +++ b/dpnp/__init__.py @@ -515,6 +515,7 @@ partition, sort, sort_complex, + top_k, ) # ----------------------------------------------------------------------------- @@ -995,6 +996,7 @@ "searchsorted", "sort", "sort_complex", + "top_k", "where", ] diff --git a/dpnp/dpnp_iface_sorting.py b/dpnp/dpnp_iface_sorting.py index 7aaf2852ec5..6ee1e75094e 100644 --- a/dpnp/dpnp_iface_sorting.py +++ b/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), + ) diff --git a/dpnp/tensor/_sorting.py b/dpnp/tensor/_sorting.py index b3aca251b04..95f25183c42 100644 --- a/dpnp/tensor/_sorting.py +++ b/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: diff --git a/dpnp/tests/tensor/test_usm_ndarray_top_k.py b/dpnp/tests/tensor/test_usm_ndarray_top_k.py index 8dce49c8592..26f491469db 100644 --- a/dpnp/tests/tensor/test_usm_ndarray_top_k.py +++ b/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): diff --git a/dpnp/tests/test_sort.py b/dpnp/tests/test_sort.py index 3834d488b8a..48bd2733074 100644 --- a/dpnp/tests/test_sort.py +++ b/dpnp/tests/test_sort.py @@ -10,6 +10,7 @@ generate_random_numpy_array, get_all_dtypes, get_complex_dtypes, + get_float_complex_dtypes, get_float_dtypes, get_integer_dtypes, ) @@ -666,3 +667,180 @@ def test_complex(self, dtype): result = dpnp.sort_complex(ia) expected = numpy.sort_complex(a) assert_equal(result, expected) + + +# TODO: compare against numpy.top_k in all tests once NumPy 2.6 is released +class TestTopK: + @staticmethod + def _expected_values(a, k, axis, mode): + res = numpy.sort(a, axis=axis, descending=(mode == "largest")) + return numpy.take(res, numpy.arange(k), axis=axis) + + def _check(self, a, ia, k, axis, mode, sorted): + values, indices = dpnp.top_k(ia, k, axis=axis, mode=mode, sorted=sorted) + + # indices of equal elements are not guaranteed to be stable, so check + # that they are distinct and point at the returned values + assert indices.dtype == dpnp.intp + np_indices = indices.asnumpy() + assert_array_equal( + numpy.take_along_axis(a, np_indices, axis=axis), values + ) + diff_axis = -1 if axis is None else axis + assert numpy.all( + numpy.diff(numpy.sort(np_indices, axis=diff_axis), axis=diff_axis) + ) + + expected = self._expected_values(a, k, axis, mode) + if not sorted: + # the order of elements is unspecified, so compare sorted results + values = dpnp.sort(values, axis=axis) + expected = numpy.sort(expected, axis=axis) + assert_array_equal(values, expected) + + @testing.with_requires("numpy>=2.5") + @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True)) + @pytest.mark.parametrize("axis", [0, 1, -1]) + @pytest.mark.parametrize("mode", ["largest", "smallest"]) + @pytest.mark.parametrize("sorted", [True, False]) + def test_basic(self, dtype, axis, mode, sorted): + a = generate_random_numpy_array((7, 30, 5), dtype) + ia = dpnp.array(a) + + for k in [0, 1, 4, a.shape[axis]]: + self._check(a, ia, k, axis, mode, sorted) + + @testing.with_requires("numpy>=2.5") + @pytest.mark.parametrize("dtype", get_float_complex_dtypes()) + @pytest.mark.parametrize("mode", ["largest", "smallest"]) + @pytest.mark.parametrize("sorted", [True, False]) + def test_nan(self, dtype, mode, sorted): + # mostly NaN, so that k <= n/2 selects NaN before sorting + a = generate_random_numpy_array((3, 100), dtype) + a[:, numpy.arange(100) % 3 != 0] = numpy.nan + ia = dpnp.array(a) + + for k in [2, 40, 50, 60, 100]: + self._check(a, ia, k, -1, mode, sorted) + + @testing.with_requires("numpy>=2.5") + @pytest.mark.parametrize("dtype", [dpnp.int16, dpnp.float32]) + @pytest.mark.parametrize("axis", [0, 1, -1]) + @pytest.mark.parametrize("mode", ["largest", "smallest"]) + @pytest.mark.parametrize("sorted", [True, False]) + def test_long_rows(self, dtype, axis, mode, sorted): + # long enough for dpnp.tensor.top_k to return the elements unsorted + a = generate_random_numpy_array((70, 80, 90), dtype) + ia = dpnp.array(a) + + n = a.shape[axis] + for k in [2, 20, n // 2, n // 2 + 1, n]: + self._check(a, ia, k, axis, mode, sorted) + + @testing.with_requires("numpy>=2.5") + @pytest.mark.parametrize( + "shape, axis, k", + [ + ((0, 5), -1, 2), + ((0, 5), 0, 0), + ((5, 0), -1, 0), + ((5, 0), 0, 3), + ((0, 5), None, 0), + ], + ) + @pytest.mark.parametrize("sorted", [True, False]) + def test_empty(self, shape, axis, k, sorted): + a = numpy.empty(shape, dtype=dpnp.int32) + ia = dpnp.array(a) + + for mode in ["largest", "smallest"]: + self._check(a, ia, k, axis, mode, sorted) + + @pytest.mark.parametrize( + "shape, axis", [((5, 0), -1), ((0, 5), 0), ((0, 0), None)] + ) + def test_empty_axis(self, shape, axis): + ia = dpnp.empty(shape, dtype=dpnp.int32) + + # numpy.top_k returns empty arrays, dpnp raises as for any k > n + with pytest.raises(ValueError, match="out of bounds"): + dpnp.top_k(ia, 1, axis=axis) + + @testing.with_requires("numpy>=2.5") + @pytest.mark.parametrize("mode", ["largest", "smallest"]) + @pytest.mark.parametrize("sorted", [True, False]) + def test_axis_none(self, mode, sorted): + a = generate_random_numpy_array((4, 5, 6), dtype=dpnp.int32) + ia = dpnp.array(a) + + for k in [0, 1, 10, a.size // 2 + 1, a.size]: + self._check(a, ia, k, None, mode, sorted) + self._check( + a[::2, :, ::-1], ia[::2, :, ::-1], k // 2, None, mode, sorted + ) + + a, ia = numpy.array(5), dpnp.array(5) + self._check(a, ia, 1, None, mode, sorted) + + def test_numpy_examples(self): + a = dpnp.array([[1, 2, 3, 4], [4, 3, 2, 1]]) + + res = dpnp.top_k(a, 2) + assert res._fields == ("values", "indices") + assert_array_equal(res.values, [[4, 3], [4, 3]]) + assert_array_equal(res.indices, [[3, 2], [0, 1]]) + + values, indices = dpnp.top_k(a, 2, axis=0) + assert_array_equal(values, [[4, 3, 3, 4], [1, 2, 2, 1]]) + assert_array_equal(indices, [[1, 1, 0, 0], [0, 0, 1, 1]]) + + values, indices = dpnp.top_k(a, 2, axis=1, mode="smallest") + assert_array_equal(values, [[1, 2], [1, 2]]) + assert_array_equal(indices, [[0, 1], [3, 2]]) + + values, indices = dpnp.top_k(dpnp.array([1.0, 2.0, 3.0, dpnp.nan]), 2) + assert_array_equal(values, [3.0, 2.0]) + assert_array_equal(indices, [2, 1]) + + @testing.with_requires("numpy>=2.6") + @pytest.mark.parametrize("axis", [0, -1]) + @pytest.mark.parametrize("mode", ["largest", "smallest"]) + def test_numpy(self, axis, mode): + a = generate_random_numpy_array((10, 20), "f4") + ia = dpnp.array(a) + + result = dpnp.top_k(ia, 5, axis=axis, mode=mode) + expected = numpy.top_k(a, 5, axis=axis, mode=mode) + assert isinstance(result, tuple) + assert_array_equal(result[0], expected[0]) + assert_array_equal(result[1], expected[1]) + + def test_usm_ndarray(self): + a = numpy.array([3, 1, 4, 1, 5, 9, 2, 6]) + ia = dpnp.tensor.asarray(a) + + values, indices = dpnp.top_k(ia, 3) + assert isinstance(values, dpnp.ndarray) + assert isinstance(indices, dpnp.ndarray) + assert_array_equal(values, [9, 6, 5]) + assert_array_equal(indices, [5, 7, 4]) + + def test_error(self): + ia = dpnp.arange(10).reshape(2, 5) + + # k is negative or out of bounds + with pytest.raises(ValueError, match="non-negative"): + dpnp.top_k(ia, -1) + assert_raises(ValueError, dpnp.top_k, ia, 6) + assert_raises(TypeError, dpnp.top_k, ia, 1.0) + + # unsupported mode + assert_raises(ValueError, dpnp.top_k, ia, 1, mode="biggest") + assert_raises(ValueError, dpnp.top_k, ia, 5, mode="biggest") + + # axis is out of bounds + assert_raises(AxisError, dpnp.top_k, ia, 1, axis=2) + assert_raises(AxisError, dpnp.top_k, dpnp.array(1), 1) + + # unsupported input type + assert_raises(TypeError, dpnp.top_k, ia.asnumpy(), 1) diff --git a/dpnp/tests/test_sycl_queue.py b/dpnp/tests/test_sycl_queue.py index e0f882f2175..40e6a9f35ad 100644 --- a/dpnp/tests/test_sycl_queue.py +++ b/dpnp/tests/test_sycl_queue.py @@ -1498,6 +1498,16 @@ def test_unique(axis, device): assert_sycl_queue_equal(iv.sycl_queue, ia.sycl_queue) +@pytest.mark.parametrize("k", [1, 2, 3]) +@pytest.mark.parametrize("sorted", [True, False]) +@pytest.mark.parametrize("device", valid_dev, ids=dev_ids) +def test_top_k(k, sorted, device): + ia = dpnp.array([[3, 1, 4, 1], [5, 9, 2, 6]], device=device) + result = dpnp.top_k(ia, k, sorted=sorted) + for iv in result: + assert_sycl_queue_equal(iv.sycl_queue, ia.sycl_queue) + + @pytest.mark.parametrize("copy", [True, False]) @pytest.mark.parametrize("device", valid_dev, ids=dev_ids) def test_nan_to_num(copy, device): diff --git a/dpnp/tests/test_usm_type.py b/dpnp/tests/test_usm_type.py index 567078c3eff..50a34fa25c4 100644 --- a/dpnp/tests/test_usm_type.py +++ b/dpnp/tests/test_usm_type.py @@ -1291,6 +1291,16 @@ def test_unique(axis, usm_type): assert x.usm_type == usm_type +@pytest.mark.parametrize("k", [1, 2, 3]) +@pytest.mark.parametrize("sorted", [True, False]) +@pytest.mark.parametrize("usm_type", list_of_usm_types) +def test_top_k(k, sorted, usm_type): + a = dpnp.array([[3, 1, 4, 1], [5, 9, 2, 6]], usm_type=usm_type) + res = dpnp.top_k(a, k, sorted=sorted) + for x in res: + assert x.usm_type == usm_type + + @pytest.mark.parametrize("copy", [True, False]) @pytest.mark.parametrize("usm_type_a", list_of_usm_types) def test_nan_to_num(copy, usm_type_a):