From d415fe229dae4230eb486bed8c47d770281db1d7 Mon Sep 17 00:00:00 2001 From: Anton <100830759+antonwolfy@users.noreply.github.com> Date: Wed, 22 Jan 2025 09:08:48 +0100 Subject: [PATCH] Add `descending` keyword argument to `dpnp.sort` and `dpnp.argsort` (#2269) The PR proposes to add `descending` keyword argument to sorting functions, including `dpnp.sort`, `dpnp.argsort` and `dpnp.ndarray.sort`, `dpnp.ndarray.argsort` methods. The keyword is mandated according to python array API. The corresponding muted tests are enabled in python array API compliance scope. --- .github/workflows/array-api-skips.txt | 6 -- dpnp/dpnp_array.py | 99 +++++++++++++++++++++++++-- dpnp/dpnp_iface_sorting.py | 73 +++++++++++++++----- dpnp/tests/test_sort.py | 59 +++++++++++++--- 4 files changed, 199 insertions(+), 38 deletions(-) diff --git a/.github/workflows/array-api-skips.txt b/.github/workflows/array-api-skips.txt index 102320283c9..f9de622d05f 100644 --- a/.github/workflows/array-api-skips.txt +++ b/.github/workflows/array-api-skips.txt @@ -47,12 +47,6 @@ array_api_tests/test_operators_and_elementwise_functions.py::test_clip array_api_tests/test_operators_and_elementwise_functions.py::test_asin array_api_tests/test_operators_and_elementwise_functions.py::test_asinh -# missing 'descending' keyword argument -array_api_tests/test_signatures.py::test_func_signature[argsort] -array_api_tests/test_signatures.py::test_func_signature[sort] -array_api_tests/test_sorting_functions.py::test_argsort -array_api_tests/test_sorting_functions.py::test_sort - # missing 'correction' keyword argument array_api_tests/test_signatures.py::test_func_signature[std] array_api_tests/test_signatures.py::test_func_signature[var] diff --git a/dpnp/dpnp_array.py b/dpnp/dpnp_array.py index 2073af9f557..ab9b29e3bad 100644 --- a/dpnp/dpnp_array.py +++ b/dpnp/dpnp_array.py @@ -700,14 +700,63 @@ def argmin(self, axis=None, out=None, *, keepdims=False): # 'argpartition', - def argsort(self, axis=-1, kind=None, order=None): + def argsort( + self, axis=-1, kind=None, order=None, *, descending=False, stable=None + ): """ Return an ndarray of indices that sort the array along the specified axis. Refer to :obj:`dpnp.argsort` for full documentation. + Parameters + ---------- + 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. + Default: ``-1``. + kind : {None, "stable", "mergesort", "radixsort"}, optional + Sorting algorithm. The default is ``None``, which uses parallel + merge-sort or parallel radix-sort algorithms depending on the array + data type. + Default: ``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). + Default: ``False``. + stable : {None, bool}, optional + Sort stability. If ``True``, the returned array will maintain the + relative order of `a` values which compare as equal. The same + behavior applies when set to ``False`` or ``None``. + Internally, this option selects ``kind="stable"``. + Default: ``None``. + + See Also + -------- + :obj:`dpnp.sort` : Return a sorted copy of an array. + :obj:`dpnp.argsort` : Return the indices that would sort an array. + :obj:`dpnp.lexsort` : Indirect stable sort on multiple keys. + :obj:`dpnp.searchsorted` : Find elements in a sorted array. + :obj:`dpnp.partition` : Partial sort. + + Examples + -------- + >>> import dpnp as np + >>> a = np.array([3, 1, 2]) + >>> a.argsort() + array([1, 2, 0]) + + >>> a = np.array([[0, 3], [2, 2]]) + >>> a.argsort(axis=0) + array([[0, 1], + [1, 0]]) + """ - return dpnp.argsort(self, axis, kind, order) + + return dpnp.argsort( + self, axis, kind, order, descending=descending, stable=stable + ) def asnumpy(self): """ @@ -1589,12 +1638,45 @@ def size(self): return self._array_obj.size - def sort(self, axis=-1, kind=None, order=None): + def sort( + self, axis=-1, kind=None, order=None, *, descending=False, stable=None + ): """ Sort an array in-place. Refer to :obj:`dpnp.sort` for full documentation. + Parameters + ---------- + axis : int, optional + Axis along which to sort. The default is ``-1``, which sorts along + the last axis. + Default: ``-1``. + kind : {None, "stable", "mergesort", "radixsort"}, optional + Sorting algorithm. The default is ``None``, which uses parallel + merge-sort or parallel radix-sort algorithms depending on the array + data type. + Default: ``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). + Default: ``False``. + stable : {None, bool}, optional + Sort stability. If ``True``, the returned array will maintain the + relative order of `a` values which compare as equal. The same + behavior applies when set to ``False`` or ``None``. + Internally, this option selects ``kind="stable"``. + Default: ``None``. + + See Also + -------- + :obj:`dpnp.sort` : Return a sorted copy of an array. + :obj:`dpnp.argsort` : Return the indices that would sort an array. + :obj:`dpnp.lexsort` : Indirect stable sort on multiple keys. + :obj:`dpnp.searchsorted` : Find elements in a sorted array. + :obj:`dpnp.partition` : Partial sort. + Note ---- `axis` in :obj:`dpnp.sort` could be integer or ``None``. If ``None``, @@ -1605,7 +1687,7 @@ def sort(self, axis=-1, kind=None, order=None): Examples -------- >>> import dpnp as np - >>> a = np.array([[1,4],[3,1]]) + >>> a = np.array([[1, 4], [3, 1]]) >>> a.sort(axis=1) >>> a array([[1, 4], @@ -1621,7 +1703,14 @@ def sort(self, axis=-1, kind=None, order=None): raise TypeError( "'NoneType' object cannot be interpreted as an integer" ) - self[...] = dpnp.sort(self, axis=axis, kind=kind, order=order) + self[...] = dpnp.sort( + self, + axis=axis, + kind=kind, + order=order, + descending=descending, + stable=stable, + ) def squeeze(self, axis=None): """ diff --git a/dpnp/dpnp_iface_sorting.py b/dpnp/dpnp_iface_sorting.py index 010491887bb..22a1f447da4 100644 --- a/dpnp/dpnp_iface_sorting.py +++ b/dpnp/dpnp_iface_sorting.py @@ -58,7 +58,13 @@ def _wrap_sort_argsort( - a, _sorting_fn, axis=-1, kind=None, order=None, stable=True + a, + _sorting_fn, + axis=-1, + kind=None, + order=None, + descending=False, + stable=True, ): """Wrap a sorting call from dpctl.tensor interface.""" @@ -83,11 +89,15 @@ 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, kind=kind) + usm_res = _sorting_fn( + usm_a, axis=axis, descending=descending, stable=stable, kind=kind + ) return dpnp_array._create_from_usm_ndarray(usm_res) -def argsort(a, axis=-1, kind=None, order=None, *, stable=None): +def argsort( + a, axis=-1, kind=None, order=None, *, descending=False, stable=None +): """ Returns the indices that would sort an array. @@ -100,13 +110,21 @@ def argsort(a, axis=-1, kind=None, order=None, *, stable=None): 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. + Default: ``-1``. kind : {None, "stable", "mergesort", "radixsort"}, optional - Sorting algorithm. Default is ``None``, which is equivalent to - ``"stable"``. + Sorting algorithm. The default is ``None``, which uses parallel + merge-sort or parallel radix-sort algorithms depending on the array + data type. + Default: ``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). + Default: ``False``. stable : {None, bool}, optional - Sort stability. If ``True``, the returned array will maintain - the relative order of ``a`` values which compare as equal. - The same behavior applies when set to ``False`` or ``None``. + Sort stability. If ``True``, the returned array will maintain the + relative order of `a` values which compare as equal. The same behavior + applies when set to ``False`` or ``None``. Internally, this option selects ``kind="stable"``. Default: ``None``. @@ -130,7 +148,6 @@ def argsort(a, axis=-1, kind=None, order=None, *, stable=None): Otherwise ``NotImplementedError`` exception will be raised. Sorting algorithms ``"quicksort"`` and ``"heapsort"`` are not supported. - See Also -------- :obj:`dpnp.ndarray.argsort` : Equivalent method. @@ -171,7 +188,13 @@ def argsort(a, axis=-1, kind=None, order=None, *, stable=None): """ return _wrap_sort_argsort( - a, dpt.argsort, axis=axis, kind=kind, order=order, stable=stable + a, + dpt.argsort, + axis=axis, + kind=kind, + order=order, + descending=descending, + stable=stable, ) @@ -215,7 +238,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=None): +def sort(a, axis=-1, kind=None, order=None, *, descending=False, stable=None): """ Return a sorted copy of an array. @@ -228,13 +251,21 @@ def sort(a, axis=-1, kind=None, order=None, *, stable=None): 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. + Default: ``-1``. kind : {None, "stable", "mergesort", "radixsort"}, optional - Sorting algorithm. Default is ``None``, which is equivalent to - ``"stable"``. + Sorting algorithm. The default is ``None``, which uses parallel + merge-sort or parallel radix-sort algorithms depending on the array + data type. + Default: ``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). + Default: ``False``. stable : {None, bool}, optional - Sort stability. If ``True``, the returned array will maintain - the relative order of ``a`` values which compare as equal. - The same behavior applies when set to ``False`` or ``None``. + Sort stability. If ``True``, the returned array will maintain the + relative order of `a` values which compare as equal. The same behavior + applies when set to ``False`` or ``None``. Internally, this option selects ``kind="stable"``. Default: ``None``. @@ -265,7 +296,7 @@ def sort(a, axis=-1, kind=None, order=None, *, stable=None): Examples -------- >>> import dpnp as np - >>> a = np.array([[1,4],[3,1]]) + >>> a = np.array([[1, 4], [3, 1]]) >>> np.sort(a) # sort along the last axis array([[1, 4], [1, 3]]) @@ -278,7 +309,13 @@ def sort(a, axis=-1, kind=None, order=None, *, stable=None): """ return _wrap_sort_argsort( - a, dpt.sort, axis=axis, kind=kind, order=order, stable=stable + a, + dpt.sort, + axis=axis, + kind=kind, + order=order, + descending=descending, + stable=stable, ) diff --git a/dpnp/tests/test_sort.py b/dpnp/tests/test_sort.py index 0a4fdfec528..5bb51311b9e 100644 --- a/dpnp/tests/test_sort.py +++ b/dpnp/tests/test_sort.py @@ -36,7 +36,7 @@ def test_axis(self, axis): result = dpnp.argsort(ia, axis=axis) expected = numpy.argsort(a, axis=axis) - assert_dtype_allclose(result, expected) + assert_array_equal(result, expected) @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True)) @pytest.mark.parametrize("axis", [None, -2, -1, 0, 1]) @@ -46,7 +46,7 @@ def test_ndarray(self, dtype, axis): result = ia.argsort(axis=axis) expected = a.argsort(axis=axis, kind="stable") - assert_dtype_allclose(result, expected) + assert_array_equal(result, expected) # this test validates that all different options of kind in dpnp are stable @pytest.mark.parametrize("kind", [None, "stable", "mergesort", "radixsort"]) @@ -56,7 +56,30 @@ def test_kind(self, kind): result = dpnp.argsort(ia, kind=kind) expected = numpy.argsort(a, kind="stable") - assert_dtype_allclose(result, expected) + assert_array_equal(result, expected) + + @pytest.mark.parametrize("descending", [False, True]) + def test_descending(self, descending): + a = numpy.repeat(numpy.arange(10), 10) + 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 + assert_array_equal(result, expected) # `stable` keyword is supported in numpy 2.0 and above @testing.with_requires("numpy>=2.0") @@ -67,7 +90,7 @@ def test_stable(self, stable): result = dpnp.argsort(ia, stable=stable) expected = numpy.argsort(a, stable=True) - assert_dtype_allclose(result, expected) + assert_array_equal(result, expected) def test_zero_dim(self): a = numpy.array(2.5) @@ -80,7 +103,7 @@ def test_zero_dim(self): # with axis = None result = dpnp.argsort(ia, axis=None) expected = numpy.argsort(a, axis=None) - assert_dtype_allclose(result, expected) + assert_array_equal(result, expected) class TestSearchSorted: @@ -273,7 +296,7 @@ def test_axis(self, axis): result = dpnp.sort(ia, axis=axis) expected = numpy.sort(a, axis=axis) - assert_dtype_allclose(result, expected) + assert_array_equal(result, expected) @pytest.mark.parametrize("dtype", get_all_dtypes()) @pytest.mark.parametrize("axis", [-2, -1, 0, 1]) @@ -293,7 +316,25 @@ def test_kind(self, kind): result = dpnp.sort(ia, kind=kind) expected = numpy.sort(a, kind="stable") - assert_dtype_allclose(result, expected) + assert_array_equal(result, expected) + + @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) + assert_array_equal(result, expected) + + # test ndarray method + ia.sort(descending=descending) + a.sort(kind="stable") + if descending: + a = numpy.flip(a) + assert_array_equal(ia, a) # `stable` keyword is supported in numpy 2.0 and above @testing.with_requires("numpy>=2.0") @@ -304,7 +345,7 @@ def test_stable(self, stable): result = dpnp.sort(ia, stable=stable) expected = numpy.sort(a, stable=True) - assert_dtype_allclose(result, expected) + assert_array_equal(result, expected) def test_ndarray_axis_none(self): a = numpy.random.uniform(-10, 10, 12) @@ -323,7 +364,7 @@ def test_zero_dim(self): # with axis = None result = dpnp.sort(ia, axis=None) expected = numpy.sort(a, axis=None) - assert_dtype_allclose(result, expected) + assert_array_equal(result, expected) def test_error(self): ia = dpnp.arange(10)