Sort NaNs to the end for descending order

CHANGELOG.md

@@ -48,6 +48,8 @@ This release is compatible with NumPy 2.5.
 * `dpnp` uses pybind11 3.1.0 [#3015](https://github.com/IntelPython/dpnp/pull/3015)
 * Reworked the ASV benchmarks and added end-to-end workload benchmarks derived from dpBench [#2996](https://github.com/IntelPython/dpnp/pull/2996)
 * Reduced allocations in `dpnp.linalg.norm` by reusing the reduction result as the `sqrt` output buffer in the 2-norm and Frobenius-norm branches [#3062](https://github.com/IntelPython/dpnp/pull/3062)
+* Changed `dpnp.sort`, `dpnp.argsort`, and their `dpnp.ndarray`/`dpnp.tensor` counterparts placing `NaN` values first instead of last when sorting in descending order [#3066](https://github.com/IntelPython/dpnp/pull/3066)
+* `dpnp.tensor.top_k` now treats `NaN` values (and complex values with a `NaN` component) as the smallest, consistent with `dpnp.tensor.sort`, so `mode="largest"` no longer returns them ahead of finite values [#3066](https://github.com/IntelPython/dpnp/pull/3066)
 
 ### Deprecated
 

dpnp/dpnp_iface_sorting.py

@@ -119,7 +119,8 @@ def argsort(
     descending : bool, optional
         Sort order. If ``True``, the array must be sorted in descending order
         (by value). If ``False``, the array must be sorted in ascending order
-        (by value).
+        (by value). NaN values (and complex values with a NaN component) are
+        ordered to the end regardless of `descending`.
 
         Default: ``False``.
     stable : {None, bool}, optional
@@ -348,7 +349,8 @@ def sort(a, axis=-1, kind=None, order=None, *, descending=False, stable=None):
     descending : bool, optional
         Sort order. If ``True``, the array must be sorted in descending order
         (by value). If ``False``, the array must be sorted in ascending order
-        (by value).
+        (by value). NaN values (and complex values with a NaN component) are
+        ordered to the end regardless of `descending`.
 
         Default: ``False``.
     stable : {None, bool}, optional

dpnp/tensor/_sorting.py

@@ -73,7 +73,9 @@ def sort(x, /, *, axis=-1, descending=False, stable=True, kind=None):
         descending (Optional[bool]):
             sort order. If `True`, the array must be sorted in descending
             order (by value). If `False`, the array must be sorted in
-            ascending order (by value). Default: `False`.
+            ascending order (by value). NaN values (and complex values with
+            a NaN component) are ordered to the end regardless of
+            `descending`. Default: `False`.
         stable (Optional[bool]):
             sort stability. If `True`, the returned array must maintain the
             relative order of `x` values which compare as equal. If `False`,
@@ -185,7 +187,9 @@ def argsort(x, axis=-1, descending=False, stable=True, kind=None):
         descending (Optional[bool]):
             sort order. If `True`, the array must be sorted in descending
             order (by value). If `False`, the array must be sorted in
-            ascending order (by value). Default: `False`.
+            ascending order (by value). NaN values (and complex values with
+            a NaN component) are ordered to the end regardless of
+            `descending`. Default: `False`.
         stable (Optional[bool]):
             sort stability. If `True`, the returned array must maintain the
             relative order of `x` values which compare as equal. If `False`,
@@ -315,6 +319,10 @@ def top_k(x, k, /, *, axis=None, mode="largest"):
             - `"largest"`: return the `k` largest elements.
             - `"smallest"`: return the `k` smallest elements.
 
+            NaN values (and complex values with a NaN component) are treated
+            as the smallest, so `"largest"` does not return them ahead of
+            finite values.
+
             Default: `"largest"`.
 
     Returns:

dpnp/tensor/libtensor/include/kernels/sorting/radix_sort.hpp

@@ -170,9 +170,11 @@ std::uint16_t order_preserving_cast(sycl::half val)
 {
     using UIntT = std::uint16_t;
 
-    const UIntT uint_val = sycl::bit_cast<UIntT>(
-        (sycl::isnan(val)) ? std::numeric_limits<sycl::half>::quiet_NaN()
-                           : val);
+    // NaNs sort to the end for both orders
+    if (sycl::isnan(val))
+        return std::numeric_limits<UIntT>::max();
+
+    const UIntT uint_val = sycl::bit_cast<UIntT>(val);
     UIntT mask;
 
     // test the sign bit of the original value
@@ -203,8 +205,11 @@ std::uint32_t order_preserving_cast(FloatT val)
 {
     using UIntT = std::uint32_t;
 
-    UIntT uint_val = sycl::bit_cast<UIntT>(
-        (sycl::isnan(val)) ? std::numeric_limits<FloatT>::quiet_NaN() : val);
+    // NaNs sort to the end for both orders
+    if (sycl::isnan(val))
+        return std::numeric_limits<UIntT>::max();
+
+    const UIntT uint_val = sycl::bit_cast<UIntT>(val);
 
     UIntT mask;
 
@@ -231,8 +236,11 @@ std::uint64_t order_preserving_cast(FloatT val)
 {
     using UIntT = std::uint64_t;
 
-    UIntT uint_val = sycl::bit_cast<UIntT>(
-        (sycl::isnan(val)) ? std::numeric_limits<FloatT>::quiet_NaN() : val);
+    // NaNs sort to the end for both orders
+    if (sycl::isnan(val))
+        return std::numeric_limits<UIntT>::max();
+
+    const UIntT uint_val = sycl::bit_cast<UIntT>(val);
     UIntT mask;
 
     // test the sign bit of the original value

dpnp/tensor/libtensor/include/utils/rich_comparisons.hpp

@@ -60,9 +60,10 @@ struct ExtendedRealFPLess
 template <typename fpT>
 struct ExtendedRealFPGreater
 {
+    /* [R, nan] — NaNs sort to the end, as in ascending order */
     bool operator()(const fpT v1, const fpT v2) const
     {
-        return (!std::isnan(v2) && (std::isnan(v1) || (v2 < v1)));
+        return (!std::isnan(v1) && (std::isnan(v2) || (v2 < v1)));
     }
 };
 
@@ -106,10 +107,38 @@ struct ExtendedComplexFPLess
 template <typename cT>
 struct ExtendedComplexFPGreater
 {
+    /* [(R, R), (R, nan), (nan, R), (nan, nan)] — NaN-containing values keep
+       the same trailing groups as ascending order; only the finite-component
+       comparison within a group is reversed */
     bool operator()(const cT &v1, const cT &v2) const
     {
-        auto less_ = ExtendedComplexFPLess<cT>{};
-        return less_(v2, v1);
+        using realT = typename cT::value_type;
+
+        const realT real1 = std::real(v1);
+        const realT real2 = std::real(v2);
+
+        const bool r1_nan = std::isnan(real1);
+        const bool r2_nan = std::isnan(real2);
+
+        const realT imag1 = std::imag(v1);
+        const realT imag2 = std::imag(v2);
+
+        const bool i1_nan = std::isnan(imag1);
+        const bool i2_nan = std::isnan(imag2);
+
+        const int idx1 = ((r1_nan) ? 2 : 0) + ((i1_nan) ? 1 : 0);
+        const int idx2 = ((r2_nan) ? 2 : 0) + ((i2_nan) ? 1 : 0);
+
+        const bool res =
+            !(r1_nan && i1_nan) &&
+            ((idx1 < idx2) ||
+             ((idx1 == idx2) &&
+              ((r1_nan && !i1_nan && (imag2 < imag1)) ||
+               (!r1_nan && i1_nan && (real2 < real1)) ||
+               (!r1_nan && !i1_nan &&
+                ((real2 < real1) || (!(real1 < real2) && (imag2 < imag1)))))));
+
+        return res;
     }
 };
 

dpnp/tests/tensor/test_usm_ndarray_sorting.py

@@ -33,6 +33,7 @@
 from numpy.testing import assert_array_equal
 
 import dpnp.tensor as dpt
+from dpnp.tests.helper import numpy_version
 
 from .helper import (
     get_queue_or_skip,
@@ -309,8 +310,9 @@ def test_sort_real_fp_nan(dtype, kind):
 
     s = dpt.sort(x, descending=True, kind=kind)
 
+    # NaNs sort to the end for descending order too matching NumPy
     expected = dpt.asarray(
-        [dpt.nan, dpt.nan, 0.2, 0.1, -0.0, 0.0, -0.1, -0.3], dtype=dtype
+        [0.2, 0.1, -0.0, 0.0, -0.1, -0.3, dpt.nan, dpt.nan], dtype=dtype
     )
 
     assert dpt.allclose(s, expected, equal_nan=True)
@@ -356,6 +358,21 @@ def test_sort_complex_fp_nan(dtype):
             r1.view(np.int64), r2.view(np.int64)
         ), f"Failed for {i} and {j}"
 
+    # complex values with a NaN component sort to the end for descending
+    # order too, matching NumPy (`descending` requires numpy>=2.5)
+    if numpy_version() >= "2.5.0":
+        s = dpt.sort(inp, descending=True)
+        expected = np.sort(dpt.asnumpy(inp), descending=True)
+        assert np.allclose(dpt.asnumpy(s), expected, equal_nan=True)
+
+        m1 = dpt.asnumpy(dpt.sort(sub_arrs, axis=1, descending=True))
+        m2 = np.sort(dpt.asnumpy(sub_arrs), axis=1, descending=True)
+        for k in range(len(pairs)):
+            i, j = pairs[k]
+            assert np.array_equal(
+                m1[k].view(np.int64), m2[k].view(np.int64)
+            ), f"Failed for {i} and {j}"
+
 
 def test_radix_sort_size_1_axis():
     get_queue_or_skip()

dpnp/tests/tensor/test_usm_ndarray_top_k.py

@@ -179,6 +179,39 @@ def test_top_k_1d_smallest(dtype, n):
     assert dpt.all(s.indices == expected_inds), (s.indices, expected_inds)
 
 
+@pytest.mark.parametrize("dtype", ["f2", "f4", "f8", "c8", "c16"])
+@pytest.mark.parametrize("mode", ["largest", "smallest"])
+def test_top_k_nan(dtype, mode):
+    # NaNs (and complex values with a NaN component) are ordered to the end
+    # for both modes, so top_k excludes them until k reaches the NaN region
+    q = get_queue_or_skip()
+    skip_if_dtype_not_supported(dtype, q)
+
+    is_complex = dtype in ("c8", "c16")
+    nan = complex(dpt.nan, dpt.nan) if is_complex else dpt.nan
+
+    def has_nan(a):
+        if is_complex:
+            return dpt.any(dpt.isnan(dpt.real(a)) | dpt.isnan(dpt.imag(a)))
+        return dpt.any(dpt.isnan(a))
+
+    # 5 distinct finite values followed by 2 NaNs, then rolled to interleave
+    x = dpt.roll(dpt.asarray([3, 1, 5, 2, 4, nan, nan], dtype=dtype), 3)
+
+    # k within the finite region: NaNs are excluded from the result
+    r = dpt.top_k(x, 3, mode=mode)
+    assert not has_nan(r.values)
+    assert dpt.all(r.values == x[r.indices])
+    extreme = [5, 4, 3] if mode == "largest" else [1, 2, 3]
+    expected = dpt.asarray(extreme, dtype=dtype)
+    assert dpt.all(dpt.sort(r.values) == dpt.sort(expected))
+
+    # k reaching into the NaN region: the 2 NaNs are ordered last
+    r = dpt.top_k(x, 7, mode=mode)
+    assert has_nan(r.values[-2:])
+    assert not has_nan(r.values[:5])
+
+
 @pytest.mark.parametrize(
     "dtype",
     [

dpnp/tests/test_sort.py

@@ -59,28 +59,26 @@ def test_kind(self, kind):
         expected = numpy.argsort(a, kind="stable")
         assert_array_equal(result, expected)
 
+    @testing.with_requires("numpy>=2.5")
     @pytest.mark.parametrize("descending", [False, True])
-    def test_descending(self, descending):
-        a = numpy.repeat(numpy.arange(10), 10)
+    @pytest.mark.parametrize(
+        "dtype", get_integer_dtypes(all_int_types=True) + [dpnp.bool]
+    )
+    def test_descending_duplicates(self, dtype, descending):
+        # a stable argsort keeps the original relative order of equal
+        # elements in both ascending and descending order
+        if dtype == dpnp.bool:
+            values = [False, True]
+        else:
+            info = numpy.iinfo(dtype)
+            values = [info.min, 1, info.max]
+        a = numpy.array(values * 2, dtype=dtype)
         ia = dpnp.array(a)
 
         result = dpnp.argsort(ia, descending=descending)
-        if not descending:
-            expected = numpy.argsort(a, kind="stable")
-        else:
-            expected = numpy.flip(numpy.argsort(numpy.flip(a), kind="stable"))
-            expected = (a.shape[0] - 1) - expected
-        assert_array_equal(result, expected)
-
-        # test ndarray method
-        result = ia.argsort(descending=descending)
-        if not descending:
-            expected = a.argsort(kind="stable")
-        else:
-            a = numpy.flip(a)
-            expected = numpy.flip(a.argsort(kind="stable"))
-            expected = (a.shape[0] - 1) - expected
+        expected = numpy.argsort(a, stable=True, descending=descending)
         assert_array_equal(result, expected)
+        assert_array_equal(dpnp.sort(ia, descending=descending), a[expected])
 
     # `stable` keyword is supported in numpy 2.0 and above
     @testing.with_requires("numpy>=2.0")
@@ -545,24 +543,57 @@ def test_kind(self, kind):
         expected = numpy.sort(a, kind="stable")
         assert_array_equal(result, expected)
 
+    @testing.with_requires("numpy>=2.5")
     @pytest.mark.parametrize("descending", [False, True])
     def test_descending(self, descending):
         a = numpy.repeat(numpy.arange(10), 10)
         ia = dpnp.array(a)
 
         result = dpnp.sort(ia, descending=descending)
-        expected = numpy.sort(a, kind="stable")
-        if descending:
-            expected = numpy.flip(expected)
+        expected = numpy.sort(a, stable=True, descending=descending)
         assert_array_equal(result, expected)
 
         # test ndarray method
         ia.sort(descending=descending)
-        a.sort(kind="stable")
-        if descending:
-            a = numpy.flip(a)
+        a.sort(stable=True, descending=descending)
         assert_array_equal(ia, a)
 
+    @testing.with_requires("numpy>=2.5")
+    @pytest.mark.parametrize("descending", [False, True])
+    @pytest.mark.parametrize("dtype", get_float_dtypes(no_float16=False))
+    def test_descending_nan(self, dtype, descending):
+        # NaNs are sorted to the end for both ascending and descending order
+        a = numpy.linspace(-50, 50, 101).astype(dtype)
+        a[::10] = numpy.nan
+        ia = dpnp.array(a)
+
+        result = dpnp.sort(ia, descending=descending)
+        expected = numpy.sort(a, stable=True, descending=descending)
+        assert_array_equal(result, expected)
+
+    @testing.with_requires("numpy>=2.5")
+    @pytest.mark.parametrize("descending", [False, True])
+    @pytest.mark.parametrize("dtype", get_complex_dtypes())
+    def test_descending_complex_nan(self, dtype, descending):
+        # NaN-containing complex values sort to the end in groups
+        # (no nan) -> (imag nan) -> (real nan) -> (all nan) for both orders;
+        # finite values keep lexicographic order (real part more significant)
+        arange = numpy.tile(numpy.arange(25), 4)
+        no_nans = arange + 1j * arange
+        im_nans = arange + complex(0, numpy.nan)
+        re_nans = complex(numpy.nan, 0) + 1j * arange
+        all_nans = numpy.full(100, complex(numpy.nan, numpy.nan))
+        a = numpy.concatenate((no_nans, im_nans, re_nans, all_nans))
+        a = a.astype(dtype)
+
+        rng = numpy.random.default_rng(0)
+        rng.shuffle(a)
+        ia = dpnp.array(a)
+
+        result = dpnp.sort(ia, descending=descending)
+        expected = numpy.sort(a, stable=True, descending=descending)
+        assert_array_equal(result, expected)
+
     # `stable` keyword is supported in numpy 2.0 and above
     @testing.with_requires("numpy>=2.0")
     @pytest.mark.parametrize("stable", [None, False, True])