diff --git a/dpnp/dpnp_iface_sorting.py b/dpnp/dpnp_iface_sorting.py index 01fc637db717..640f17da55b9 100644 --- a/dpnp/dpnp_iface_sorting.py +++ b/dpnp/dpnp_iface_sorting.py @@ -66,9 +66,10 @@ def _wrap_sort_argsort( raise NotImplementedError( "order keyword argument is only supported with its default value." ) - if kind is not None and kind != "stable": - raise NotImplementedError( - "kind keyword argument can only be None or 'stable'." + if kind is not None and stable is not None: + raise ValueError( + "`kind` and `stable` parameters can't be provided at the same time." + " Use only one of them." ) usm_a = dpnp.get_usm_ndarray(a) @@ -77,11 +78,11 @@ def _wrap_sort_argsort( axis = -1 axis = normalize_axis_index(axis, ndim=usm_a.ndim) - usm_res = _sorting_fn(usm_a, axis=axis, stable=stable) + usm_res = _sorting_fn(usm_a, axis=axis, stable=stable, kind=kind) return dpnp_array._create_from_usm_ndarray(usm_res) -def argsort(a, axis=-1, kind=None, order=None, *, stable=True): +def argsort(a, axis=-1, kind=None, order=None, *, stable=None): """ Returns the indices that would sort an array. @@ -94,9 +95,9 @@ def argsort(a, axis=-1, kind=None, order=None, *, stable=True): axis : {None, int}, optional Axis along which to sort. If ``None``, the array is flattened before sorting. The default is ``-1``, which sorts along the last axis. - kind : {None, "stable"}, optional + kind : {None, "stable", "mergesort", "radixsort"}, optional Sorting algorithm. Default is ``None``, which is equivalent to - ``"stable"``. Unlike NumPy, no other option is accepted here. + ``"stable"``. stable : {None, bool}, optional Sort stability. If ``True``, the returned array will maintain the relative order of ``a`` values which compare as equal. @@ -121,8 +122,9 @@ def argsort(a, axis=-1, kind=None, order=None, *, stable=True): Limitations ----------- Parameters `order` is only supported with its default value. - Parameter `kind` can only be ``None`` or ``"stable"`` which are equivalent. Otherwise ``NotImplementedError`` exception will be raised. + Sorting algorithms ``"quicksort"`` and ``"heapsort"`` are not supported. + See Also -------- @@ -203,7 +205,7 @@ def partition(x1, kth, axis=-1, kind="introselect", order=None): return call_origin(numpy.partition, x1, kth, axis, kind, order) -def sort(a, axis=-1, kind=None, order=None, *, stable=True): +def sort(a, axis=-1, kind=None, order=None, *, stable=None): """ Return a sorted copy of an array. @@ -216,9 +218,9 @@ def sort(a, axis=-1, kind=None, order=None, *, stable=True): axis : {None, int}, optional Axis along which to sort. If ``None``, the array is flattened before sorting. The default is ``-1``, which sorts along the last axis. - kind : {None, "stable"}, optional + kind : {None, "stable", "mergesort", "radixsort"}, optional Sorting algorithm. Default is ``None``, which is equivalent to - ``"stable"``. Unlike NumPy, no other option is accepted here. + ``"stable"``. stable : {None, bool}, optional Sort stability. If ``True``, the returned array will maintain the relative order of ``a`` values which compare as equal. @@ -239,8 +241,8 @@ def sort(a, axis=-1, kind=None, order=None, *, stable=True): Limitations ----------- Parameters `order` is only supported with its default value. - Parameter `kind` can only be ``None`` or ``"stable"`` which are equivalent. Otherwise ``NotImplementedError`` exception will be raised. + Sorting algorithms ``"quicksort"`` and ``"heapsort"`` are not supported. See Also -------- diff --git a/tests/test_sort.py b/tests/test_sort.py index afaf4e152e7a..07156b5d882a 100644 --- a/tests/test_sort.py +++ b/tests/test_sort.py @@ -15,29 +15,34 @@ class TestArgsort: + @pytest.mark.parametrize("kind", [None, "stable", "mergesort", "radixsort"]) @pytest.mark.parametrize("dtype", get_all_dtypes(no_complex=True)) - def test_argsort_dtype(self, dtype): + def test_basic(self, kind, dtype): a = numpy.random.uniform(-5, 5, 10) np_array = numpy.array(a, dtype=dtype) dp_array = dpnp.array(np_array) - result = dpnp.argsort(dp_array, kind="stable") + result = dpnp.argsort(dp_array, kind=kind) expected = numpy.argsort(np_array, kind="stable") assert_dtype_allclose(result, expected) + @pytest.mark.parametrize("kind", [None, "stable", "mergesort", "radixsort"]) @pytest.mark.parametrize("dtype", get_complex_dtypes()) - def test_argsort_complex(self, dtype): + def test_complex(self, kind, dtype): a = numpy.random.uniform(-5, 5, 10) b = numpy.random.uniform(-5, 5, 10) np_array = numpy.array(a + b * 1j, dtype=dtype) dp_array = dpnp.array(np_array) - result = dpnp.argsort(dp_array) - expected = numpy.argsort(np_array) - assert_dtype_allclose(result, expected) + if kind == "radixsort": + assert_raises(ValueError, dpnp.argsort, dp_array, kind=kind) + else: + result = dpnp.argsort(dp_array, kind=kind) + expected = numpy.argsort(np_array) + assert_dtype_allclose(result, expected) @pytest.mark.parametrize("axis", [None, -2, -1, 0, 1, 2]) - def test_argsort_axis(self, axis): + def test_axis(self, axis): a = numpy.random.uniform(-10, 10, 36) np_array = numpy.array(a).reshape(3, 4, 3) dp_array = dpnp.array(np_array) @@ -48,7 +53,7 @@ def test_argsort_axis(self, axis): @pytest.mark.parametrize("dtype", get_all_dtypes()) @pytest.mark.parametrize("axis", [None, -2, -1, 0, 1]) - def test_argsort_ndarray(self, dtype, axis): + def test_ndarray(self, dtype, axis): if dtype and issubclass(dtype, numpy.integer): a = numpy.random.choice( numpy.arange(-10, 10), replace=False, size=12 @@ -62,8 +67,9 @@ def test_argsort_ndarray(self, dtype, axis): expected = np_array.argsort(axis=axis) assert_dtype_allclose(result, expected) - @pytest.mark.parametrize("kind", [None, "stable"]) - def test_sort_kind(self, kind): + # this test validates that all different options of kind in dpnp are stable + @pytest.mark.parametrize("kind", [None, "stable", "mergesort", "radixsort"]) + def test_kind(self, kind): np_array = numpy.repeat(numpy.arange(10), 10) dp_array = dpnp.array(np_array) @@ -74,7 +80,7 @@ def test_sort_kind(self, kind): # `stable` keyword is supported in numpy 2.0 and above @testing.with_requires("numpy>=2.0") @pytest.mark.parametrize("stable", [None, False, True]) - def test_sort_stable(self, stable): + def test_stable(self, stable): np_array = numpy.repeat(numpy.arange(10), 10) dp_array = dpnp.array(np_array) @@ -82,7 +88,7 @@ def test_sort_stable(self, stable): expected = numpy.argsort(np_array, stable=True) assert_dtype_allclose(result, expected) - def test_argsort_zero_dim(self): + def test_zero_dim(self): np_array = numpy.array(2.5) dp_array = dpnp.array(np_array) @@ -266,29 +272,34 @@ def test_v_scalar(self): class TestSort: + @pytest.mark.parametrize("kind", [None, "stable", "mergesort", "radixsort"]) @pytest.mark.parametrize("dtype", get_all_dtypes(no_complex=True)) - def test_sort_dtype(self, dtype): + def test_basic(self, kind, dtype): a = numpy.random.uniform(-5, 5, 10) np_array = numpy.array(a, dtype=dtype) dp_array = dpnp.array(np_array) - result = dpnp.sort(dp_array) + result = dpnp.sort(dp_array, kind=kind) expected = numpy.sort(np_array) assert_dtype_allclose(result, expected) + @pytest.mark.parametrize("kind", [None, "stable", "mergesort", "radixsort"]) @pytest.mark.parametrize("dtype", get_complex_dtypes()) - def test_sort_complex(self, dtype): + def test_complex(self, kind, dtype): a = numpy.random.uniform(-5, 5, 10) b = numpy.random.uniform(-5, 5, 10) np_array = numpy.array(a + b * 1j, dtype=dtype) dp_array = dpnp.array(np_array) - result = dpnp.sort(dp_array) - expected = numpy.sort(np_array) - assert_dtype_allclose(result, expected) + if kind == "radixsort": + assert_raises(ValueError, dpnp.argsort, dp_array, kind=kind) + else: + result = dpnp.sort(dp_array, kind=kind) + expected = numpy.sort(np_array) + assert_dtype_allclose(result, expected) @pytest.mark.parametrize("axis", [None, -2, -1, 0, 1, 2]) - def test_sort_axis(self, axis): + def test_axis(self, axis): a = numpy.random.uniform(-10, 10, 36) np_array = numpy.array(a).reshape(3, 4, 3) dp_array = dpnp.array(np_array) @@ -299,7 +310,7 @@ def test_sort_axis(self, axis): @pytest.mark.parametrize("dtype", get_all_dtypes()) @pytest.mark.parametrize("axis", [-2, -1, 0, 1]) - def test_sort_ndarray(self, dtype, axis): + def test_ndarray(self, dtype, axis): a = numpy.random.uniform(-10, 10, 12) np_array = numpy.array(a, dtype=dtype).reshape(6, 2) dp_array = dpnp.array(np_array) @@ -308,8 +319,9 @@ def test_sort_ndarray(self, dtype, axis): np_array.sort(axis=axis) assert_dtype_allclose(dp_array, np_array) - @pytest.mark.parametrize("kind", [None, "stable"]) - def test_sort_kind(self, kind): + # this test validates that all different options of kind in dpnp are stable + @pytest.mark.parametrize("kind", [None, "stable", "mergesort", "radixsort"]) + def test_kind(self, kind): np_array = numpy.repeat(numpy.arange(10), 10) dp_array = dpnp.array(np_array) @@ -320,7 +332,7 @@ def test_sort_kind(self, kind): # `stable` keyword is supported in numpy 2.0 and above @testing.with_requires("numpy>=2.0") @pytest.mark.parametrize("stable", [None, False, True]) - def test_sort_stable(self, stable): + def test_stable(self, stable): np_array = numpy.repeat(numpy.arange(10), 10) dp_array = dpnp.array(np_array) @@ -328,13 +340,13 @@ def test_sort_stable(self, stable): expected = numpy.sort(np_array, stable=True) assert_dtype_allclose(result, expected) - def test_sort_ndarray_axis_none(self): + def test_ndarray_axis_none(self): a = numpy.random.uniform(-10, 10, 12) dp_array = dpnp.array(a).reshape(6, 2) with pytest.raises(TypeError): dp_array.sort(axis=None) - def test_sort_zero_dim(self): + def test_zero_dim(self): np_array = numpy.array(2.5) dp_array = dpnp.array(np_array) @@ -347,15 +359,20 @@ def test_sort_zero_dim(self): expected = numpy.sort(np_array, axis=None) assert_dtype_allclose(result, expected) - def test_sort_notimplemented(self): + def test_error(self): dp_array = dpnp.arange(10) - with pytest.raises(NotImplementedError): + # quicksort is currently not supported + with pytest.raises(ValueError): dpnp.sort(dp_array, kind="quicksort") with pytest.raises(NotImplementedError): dpnp.sort(dp_array, order=["age"]) + # both kind and stable are given + with pytest.raises(ValueError): + dpnp.sort(dp_array, kind="mergesort", stable=True) + class TestSortComplex: @pytest.mark.parametrize(