From ad4e9bf6f1f80715131705ea1a6f8d54fdefad84 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Fri, 30 Oct 2026 00:45:16 -0700
Subject: [PATCH 1/5] Add dpnp.top_k

---
 doc/reference/sort.rst     |   1 +
 dpnp/__init__.py           |   2 +
 dpnp/dpnp_iface_sorting.py | 104 +++++++++++++++++++++++++++++++++++
 dpnp/tests/test_sort.py    | 109 +++++++++++++++++++++++++++++++++++++
 4 files changed, 216 insertions(+)

diff --git a/doc/reference/sort.rst b/doc/reference/sort.rst
index 2f73aa775406..0711f8e5485d 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 68a683d8095d..ce3b196e070d 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 7aaf2852ec57..9d8e1c11683a 100644
--- a/dpnp/dpnp_iface_sorting.py
+++ b/dpnp/dpnp_iface_sorting.py
@@ -448,3 +448,107 @@ 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 : int, optional
+        Axis along which to find the largest or smallest elements.
+
+        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
+    -------
+    out : tuple of dpnp.ndarray
+        A tuple ``(values, indices)``, where `values` are the top `k` values
+        of `a` and `indices` are their indices along `axis`. Both arrays have
+        the shape of `a` with the dimension along `axis` replaced by `k`.
+        `values` has the same data type as `a` and `indices` has the default
+        array index data type.
+
+    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`.
+
+    Limitations
+    -----------
+    Parameter `axis` is not supported as ``None``, same as in NumPy.
+    Otherwise ``ValueError`` exception will be raised.
+
+    See Also
+    --------
+    :obj:`dpnp.tensor.top_k` : Equivalent function in :obj:`dpnp.tensor`.
+    :obj:`dpnp.partition` : Partial sort.
+    :obj:`dpnp.sort` : Full sorting.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> a = np.array([[1, 2, 3, 4], [4, 3, 2, 1]])
+    >>> np.top_k(a, 2)
+    (array([[4, 3],
+           [4, 3]]), array([[3, 2],
+           [0, 1]]))
+    >>> np.top_k(a, 2, axis=0)
+    (array([[4, 3, 3, 4],
+           [1, 2, 2, 1]]), array([[1, 1, 0, 0],
+           [0, 0, 1, 1]]))
+    >>> np.top_k(a, 2, axis=1, mode="smallest")
+    (array([[1, 2],
+           [1, 2]]), array([[0, 1],
+           [3, 2]]))
+    >>> np.top_k(np.array([1.0, 2.0, 3.0, np.nan]), 2)
+    (array([3., 2.]), array([2, 1]))
+
+    """
+
+    if axis is None:
+        raise ValueError(
+            "`axis=None` is not supported. Please provide a valid axis."
+        )
+
+    usm_a = dpnp.get_usm_ndarray(a)
+    axis = normalize_axis_index(axis, ndim=usm_a.ndim)
+    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 usm_vals.shape[axis] > 1:
+        sort_inds = dpt.argsort(
+            usm_vals, axis=axis, descending=(mode == "largest")
+        )
+        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 (
+        dpnp_array._create_from_usm_ndarray(usm_vals),
+        dpnp_array._create_from_usm_ndarray(usm_inds),
+    )
diff --git a/dpnp/tests/test_sort.py b/dpnp/tests/test_sort.py
index 3834d488b8a7..1a4b072a5262 100644
--- a/dpnp/tests/test_sort.py
+++ b/dpnp/tests/test_sort.py
@@ -666,3 +666,112 @@ def test_complex(self, dtype):
         result = dpnp.sort_complex(ia)
         expected = numpy.sort_complex(a)
         assert_equal(result, expected)
+
+
+@testing.with_requires("numpy>=2.5")
+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)
+        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)
+
+        # indices of equal elements are not guaranteed to be stable
+        assert indices.dtype == dpnp.intp
+        assert_array_equal(
+            numpy.take_along_axis(a, indices.asnumpy(), axis=axis),
+            values if sorted else dpnp.take_along_axis(ia, indices, axis=axis),
+        )
+
+    @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)
+
+    @pytest.mark.parametrize("dtype", get_float_dtypes() + get_complex_dtypes())
+    @pytest.mark.parametrize("mode", ["largest", "smallest"])
+    @pytest.mark.parametrize("sorted", [True, False])
+    def test_nan(self, dtype, mode, sorted):
+        a = generate_random_numpy_array((3, 10), dtype)
+        a[:, ::3] = numpy.nan
+        ia = dpnp.array(a)
+
+        for k in [2, 6, 10]:
+            self._check(a, ia, k, -1, mode, sorted)
+
+    def test_numpy_examples(self):
+        a = dpnp.array([[1, 2, 3, 4], [4, 3, 2, 1]])
+
+        values, indices = dpnp.top_k(a, 2)
+        assert_array_equal(values, [[4, 3], [4, 3]])
+        assert_array_equal(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)
+
+        # axis=None is not supported
+        assert_raises(ValueError, dpnp.top_k, ia, 1, axis=None)
+
+        # k is negative or out of bounds
+        assert_raises(ValueError, 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")
+
+        # 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)

From 9cf40542f5d2d054683897af9340b0eb383216e5 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Fri, 30 Oct 2026 00:47:31 -0700
Subject: [PATCH 2/5] Add #3133 to the changelog

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b4771ac5ae86..1a3f53902fbf 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
 

From 9bfe2942b21695c714e2b20b21d03020793fa5d1 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Sun, 1 Nov 2026 18:32:04 -0800
Subject: [PATCH 3/5] Address review feedback

---
 dpnp/dpnp_iface_sorting.py    | 51 ++++++++++++++++++-------------
 dpnp/tests/test_sort.py       | 56 ++++++++++++++++++++++++++++-------
 dpnp/tests/test_sycl_queue.py | 10 +++++++
 dpnp/tests/test_usm_type.py   | 10 +++++++
 4 files changed, 96 insertions(+), 31 deletions(-)

diff --git a/dpnp/dpnp_iface_sorting.py b/dpnp/dpnp_iface_sorting.py
index 9d8e1c11683a..00a5a1f5abc9 100644
--- a/dpnp/dpnp_iface_sorting.py
+++ b/dpnp/dpnp_iface_sorting.py
@@ -39,6 +39,7 @@
 
 """
 
+import operator
 from collections.abc import Sequence
 
 import dpnp
@@ -464,8 +465,9 @@ def top_k(a, k, /, *, axis=-1, mode="largest", sorted=True):
     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 : int, optional
-        Axis along which to find the largest or smallest elements.
+    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
@@ -499,14 +501,10 @@ def top_k(a, k, /, *, axis=-1, mode="largest", sorted=True):
     the indices of equal values is not guaranteed to match their order in
     `a`. This is the case regardless of the value of `sorted`.
 
-    Limitations
-    -----------
-    Parameter `axis` is not supported as ``None``, same as in NumPy.
-    Otherwise ``ValueError`` exception will be raised.
-
     See Also
     --------
-    :obj:`dpnp.tensor.top_k` : Equivalent function in :obj:`dpnp.tensor`.
+    :obj:`dpnp.tensor.top_k` : Strictly Array API compatible version; flattens
+        by default and returns the elements in an unspecified order.
     :obj:`dpnp.partition` : Partial sort.
     :obj:`dpnp.sort` : Full sorting.
 
@@ -531,22 +529,35 @@ def top_k(a, k, /, *, axis=-1, mode="largest", sorted=True):
 
     """
 
-    if axis is None:
+    if mode not in ("largest", "smallest"):
         raise ValueError(
-            "`axis=None` is not supported. Please provide a valid axis."
+            f'mode(="{mode}") must be either "largest" or "smallest".'
         )
 
     usm_a = dpnp.get_usm_ndarray(a)
-    axis = normalize_axis_index(axis, ndim=usm_a.ndim)
-    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 usm_vals.shape[axis] > 1:
-        sort_inds = dpt.argsort(
-            usm_vals, axis=axis, descending=(mode == "largest")
-        )
-        usm_vals = dpt.take_along_axis(usm_vals, sort_inds, axis=axis)
-        usm_inds = dpt.take_along_axis(usm_inds, sort_inds, axis=axis)
+    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 (
         dpnp_array._create_from_usm_ndarray(usm_vals),
diff --git a/dpnp/tests/test_sort.py b/dpnp/tests/test_sort.py
index 1a4b072a5262..8a417e7f3b42 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,
 )
@@ -668,6 +669,7 @@ def test_complex(self, dtype):
         assert_equal(result, expected)
 
 
+# TODO: compare against numpy.top_k in all tests once NumPy 2.6 is released
 @testing.with_requires("numpy>=2.5")
 class TestTopK:
     @staticmethod
@@ -677,6 +679,19 @@ def _expected_values(a, k, axis, mode):
 
     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
@@ -684,13 +699,6 @@ def _check(self, a, ia, k, axis, mode, sorted):
             expected = numpy.sort(expected, axis=axis)
         assert_array_equal(values, expected)
 
-        # indices of equal elements are not guaranteed to be stable
-        assert indices.dtype == dpnp.intp
-        assert_array_equal(
-            numpy.take_along_axis(a, indices.asnumpy(), axis=axis),
-            values if sorted else dpnp.take_along_axis(ia, indices, axis=axis),
-        )
-
     @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True))
     @pytest.mark.parametrize("axis", [0, 1, -1])
     @pytest.mark.parametrize("mode", ["largest", "smallest"])
@@ -702,7 +710,7 @@ def test_basic(self, dtype, axis, mode, sorted):
         for k in [0, 1, 4, a.shape[axis]]:
             self._check(a, ia, k, axis, mode, sorted)
 
-    @pytest.mark.parametrize("dtype", get_float_dtypes() + get_complex_dtypes())
+    @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):
@@ -713,6 +721,34 @@ def test_nan(self, dtype, mode, sorted):
         for k in [2, 6, 10]:
             self._check(a, ia, k, -1, mode, sorted)
 
+    @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((300, 1100), dtype)
+        ia = dpnp.array(a)
+
+        n = a.shape[axis]
+        for k in [2, 100, n // 2, n // 2 + 1, n]:
+            self._check(a, ia, k, axis, mode, sorted)
+
+    @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]])
 
@@ -758,9 +794,6 @@ def test_usm_ndarray(self):
     def test_error(self):
         ia = dpnp.arange(10).reshape(2, 5)
 
-        # axis=None is not supported
-        assert_raises(ValueError, dpnp.top_k, ia, 1, axis=None)
-
         # k is negative or out of bounds
         assert_raises(ValueError, dpnp.top_k, ia, -1)
         assert_raises(ValueError, dpnp.top_k, ia, 6)
@@ -768,6 +801,7 @@ def test_error(self):
 
         # 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)
diff --git a/dpnp/tests/test_sycl_queue.py b/dpnp/tests/test_sycl_queue.py
index e0f882f2175d..9771be4e84b3 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, 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 567078c3eff2..e4184e171eaf 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, 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):

From 7bbe9f52c13a9d834ab99714ed280540e5ee518e Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 2 Nov 2026 21:49:49 -0800
Subject: [PATCH 4/5] Address review feedback

Return a `TopKResult` namedtuple from `dpnp.top_k`, document that an
empty axis with `k > 0` raises, and say "non-negative" in the
`dpnp.tensor.top_k` error for a negative `k`.

Make the tests reach the re-sort path with NaN, along every axis of a
3-D input and in the queue and USM type tests, add empty-input tests
and keep the tests that don't need NumPy 2.5 running on older NumPy.
---
 dpnp/dpnp_iface_sorting.py                  | 80 ++++++++++++++-------
 dpnp/tensor/_sorting.py                     |  4 +-
 dpnp/tests/tensor/test_usm_ndarray_top_k.py |  3 +-
 dpnp/tests/test_sort.py                     | 55 +++++++++++---
 dpnp/tests/test_sycl_queue.py               |  2 +-
 dpnp/tests/test_usm_type.py                 |  2 +-
 6 files changed, 103 insertions(+), 43 deletions(-)

diff --git a/dpnp/dpnp_iface_sorting.py b/dpnp/dpnp_iface_sorting.py
index 00a5a1f5abc9..6ee1e75094ed 100644
--- a/dpnp/dpnp_iface_sorting.py
+++ b/dpnp/dpnp_iface_sorting.py
@@ -41,6 +41,7 @@
 
 import operator
 from collections.abc import Sequence
+from typing import NamedTuple
 
 import dpnp
 import dpnp.tensor as dpt
@@ -56,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,
@@ -488,12 +495,19 @@ def top_k(a, k, /, *, axis=-1, mode="largest", sorted=True):
 
     Returns
     -------
-    out : tuple of dpnp.ndarray
-        A tuple ``(values, indices)``, where `values` are the top `k` values
-        of `a` and `indices` are their indices along `axis`. Both arrays have
-        the shape of `a` with the dimension along `axis` replaced by `k`.
-        `values` has the same data type as `a` and `indices` has the default
-        array index data type.
+    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
     -----
@@ -501,31 +515,43 @@ def top_k(a, k, /, *, axis=-1, mode="largest", sorted=True):
     the indices of equal values is not guaranteed to match their order in
     `a`. This is the case regardless of the value of `sorted`.
 
-    See Also
-    --------
-    :obj:`dpnp.tensor.top_k` : Strictly Array API compatible version; flattens
-        by default and returns the elements in an unspecified order.
-    :obj:`dpnp.partition` : Partial sort.
-    :obj:`dpnp.sort` : Full sorting.
+    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]])
-    >>> np.top_k(a, 2)
-    (array([[4, 3],
-           [4, 3]]), array([[3, 2],
-           [0, 1]]))
-    >>> np.top_k(a, 2, axis=0)
-    (array([[4, 3, 3, 4],
-           [1, 2, 2, 1]]), array([[1, 1, 0, 0],
-           [0, 0, 1, 1]]))
-    >>> np.top_k(a, 2, axis=1, mode="smallest")
-    (array([[1, 2],
-           [1, 2]]), array([[0, 1],
-           [3, 2]]))
-    >>> np.top_k(np.array([1.0, 2.0, 3.0, np.nan]), 2)
-    (array([3., 2.]), array([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])
 
     """
 
@@ -559,7 +585,7 @@ def top_k(a, k, /, *, axis=-1, mode="largest", sorted=True):
             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 (
+    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 131bb358a2bf..30d05e23df0e 100644
--- a/dpnp/tensor/_sorting.py
+++ b/dpnp/tensor/_sorting.py
@@ -309,7 +309,7 @@ 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``.
@@ -350,7 +350,7 @@ 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 None:
diff --git a/dpnp/tests/tensor/test_usm_ndarray_top_k.py b/dpnp/tests/tensor/test_usm_ndarray_top_k.py
index 8dce49c85928..5dde4cbdab03 100644
--- a/dpnp/tests/tensor/test_usm_ndarray_top_k.py
+++ b/dpnp/tests/tensor/test_usm_ndarray_top_k.py
@@ -477,8 +477,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
diff --git a/dpnp/tests/test_sort.py b/dpnp/tests/test_sort.py
index 8a417e7f3b42..48bd2733074d 100644
--- a/dpnp/tests/test_sort.py
+++ b/dpnp/tests/test_sort.py
@@ -670,7 +670,6 @@ def test_complex(self, dtype):
 
 
 # TODO: compare against numpy.top_k in all tests once NumPy 2.6 is released
-@testing.with_requires("numpy>=2.5")
 class TestTopK:
     @staticmethod
     def _expected_values(a, k, axis, mode):
@@ -699,6 +698,7 @@ def _check(self, a, ia, k, axis, mode, sorted):
             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"])
@@ -710,30 +710,63 @@ def test_basic(self, dtype, axis, mode, sorted):
         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):
-        a = generate_random_numpy_array((3, 10), dtype)
-        a[:, ::3] = numpy.nan
+        # 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, 6, 10]:
+        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((300, 1100), dtype)
+        a = generate_random_numpy_array((70, 80, 90), dtype)
         ia = dpnp.array(a)
 
         n = a.shape[axis]
-        for k in [2, 100, n // 2, n // 2 + 1, n]:
+        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):
@@ -752,9 +785,10 @@ def test_axis_none(self, mode, sorted):
     def test_numpy_examples(self):
         a = dpnp.array([[1, 2, 3, 4], [4, 3, 2, 1]])
 
-        values, indices = dpnp.top_k(a, 2)
-        assert_array_equal(values, [[4, 3], [4, 3]])
-        assert_array_equal(indices, [[3, 2], [0, 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]])
@@ -795,7 +829,8 @@ def test_error(self):
         ia = dpnp.arange(10).reshape(2, 5)
 
         # k is negative or out of bounds
-        assert_raises(ValueError, dpnp.top_k, ia, -1)
+        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)
 
diff --git a/dpnp/tests/test_sycl_queue.py b/dpnp/tests/test_sycl_queue.py
index 9771be4e84b3..40e6a9f35ad7 100644
--- a/dpnp/tests/test_sycl_queue.py
+++ b/dpnp/tests/test_sycl_queue.py
@@ -1498,7 +1498,7 @@ def test_unique(axis, device):
         assert_sycl_queue_equal(iv.sycl_queue, ia.sycl_queue)
 
 
-@pytest.mark.parametrize("k", [1, 3])
+@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):
diff --git a/dpnp/tests/test_usm_type.py b/dpnp/tests/test_usm_type.py
index e4184e171eaf..50a34fa25c4a 100644
--- a/dpnp/tests/test_usm_type.py
+++ b/dpnp/tests/test_usm_type.py
@@ -1291,7 +1291,7 @@ def test_unique(axis, usm_type):
         assert x.usm_type == usm_type
 
 
-@pytest.mark.parametrize("k", [1, 3])
+@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):

From 5d4465a090c0043ac0a8e7fdaf734799d18c7241 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 2 Nov 2026 21:49:49 -0800
Subject: [PATCH 5/5] Deprecate the flattening default of `dpnp.tensor.top_k`

`dpnp.tensor.top_k` defaults to `axis=None` and searches the flattened
array, as in the first revision of data-apis/array-api#722. The default
was changed to `axis=-1` before that PR was merged, and the 2026.12
draft of the standard has `axis=-1`.

Omitting `axis` for an input that is not 1-D now emits a
`DeprecationWarning` and keeps flattening. An explicit `axis=None` and
1-D input, for which both defaults give the same result, don't warn.
---
 CHANGELOG.md                                |  1 +
 dpnp/tensor/_sorting.py                     | 25 +++++++++++++++-
 dpnp/tests/tensor/test_usm_ndarray_top_k.py | 33 +++++++++++++++++++--
 3 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 154c7ac9373b..c065e9a29729 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -20,6 +20,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/dpnp/tensor/_sorting.py b/dpnp/tensor/_sorting.py
index 30d05e23df0e..86f9fc916bce 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
@@ -313,6 +318,12 @@ def top_k(x, k, /, *, axis=None, mode="largest"):
         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:
 
@@ -353,6 +364,18 @@ def top_k(x, k, /, *, axis=None, mode="largest"):
         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 5dde4cbdab03..26f491469db0 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)
 
@@ -489,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):