diff --git a/CHANGELOG.md b/CHANGELOG.md index b2e11822431a..e98a6763c4a4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -49,6 +49,7 @@ Also, that release drops support for Python 3.9, making Python 3.10 the minimum * Aligned `dpnp.trim_zeros` with NumPy 2.4 to support a tuple of integers passed with `axis` keyword [#2746](https://github.com/IntelPython/dpnp/pull/2746) * Aligned `strides` property of `dpnp.ndarray` with NumPy and CuPy implementations [#2747](https://github.com/IntelPython/dpnp/pull/2747) * Extended `dpnp.nan_to_num` to support broadcasting of `nan`, `posinf`, and `neginf` keywords [#2754](https://github.com/IntelPython/dpnp/pull/2754) +* Changed `dpnp.partition` implementation to reuse `dpnp.sort` where it brings the performance benefit [#2766](https://github.com/IntelPython/dpnp/pull/2766) ### Deprecated diff --git a/dpnp/backend/kernels/dpnp_krnl_sorting.cpp b/dpnp/backend/kernels/dpnp_krnl_sorting.cpp index 8a7ee9c8418f..86901e32df08 100644 --- a/dpnp/backend/kernels/dpnp_krnl_sorting.cpp +++ b/dpnp/backend/kernels/dpnp_krnl_sorting.cpp @@ -70,90 +70,25 @@ DPCTLSyclEventRef dpnp_partition_c(DPCTLSyclQueueRef q_ref, sycl::queue q = *(reinterpret_cast(q_ref)); - if (ndim == 1) // 1d array with C-contiguous data - { - _DataType *arr = static_cast<_DataType *>(array1_in); - _DataType *result = static_cast<_DataType *>(result1); + _DataType *arr = static_cast<_DataType *>(array1_in); + _DataType *result = static_cast<_DataType *>(result1); - auto policy = oneapi::dpl::execution::make_device_policy< - dpnp_partition_c_kernel<_DataType>>(q); + auto policy = oneapi::dpl::execution::make_device_policy< + dpnp_partition_c_kernel<_DataType>>(q); - // fill the result array with data from input one - q.memcpy(result, arr, size * sizeof(_DataType)).wait(); + // fill the result array with data from input one + q.memcpy(result, arr, size * sizeof(_DataType)).wait(); - // make a partial sorting such that: + for (size_t i = 0; i < size_; i++) { + _DataType *bufptr = result + i * shape_[0]; + + // for every slice it makes a partial sorting such that: // 1. result[0 <= i < kth] <= result[kth] // 2. result[kth <= i < size] >= result[kth] // event-blocking call, no need for wait() - std::nth_element(policy, result, result + kth, result + size, + std::nth_element(policy, bufptr, bufptr + kth, bufptr + size, dpnp_less_comp()); - return event_ref; - } - - DPNPC_ptr_adapter<_DataType> input1_ptr(q_ref, array1_in, size, true); - DPNPC_ptr_adapter<_DataType> input2_ptr(q_ref, array2_in, size, true); - DPNPC_ptr_adapter<_DataType> result1_ptr(q_ref, result1, size, true, true); - _DataType *arr = input1_ptr.get_ptr(); - _DataType *arr2 = input2_ptr.get_ptr(); - _DataType *result = result1_ptr.get_ptr(); - - auto arr_to_result_event = q.memcpy(result, arr, size * sizeof(_DataType)); - arr_to_result_event.wait(); - - _DataType *matrix = new _DataType[shape_[ndim - 1]]; - - for (size_t i = 0; i < size_; ++i) { - size_t ind_begin = i * shape_[ndim - 1]; - size_t ind_end = (i + 1) * shape_[ndim - 1] - 1; - - for (size_t j = ind_begin; j < ind_end + 1; ++j) { - size_t ind = j - ind_begin; - matrix[ind] = arr2[j]; - } - std::partial_sort(matrix, matrix + shape_[ndim - 1], - matrix + shape_[ndim - 1], dpnp_less_comp()); - for (size_t j = ind_begin; j < ind_end + 1; ++j) { - size_t ind = j - ind_begin; - arr2[j] = matrix[ind]; - } } - - shape_elem_type *shape = reinterpret_cast( - sycl::malloc_shared(ndim * sizeof(shape_elem_type), q)); - auto memcpy_event = q.memcpy(shape, shape_, ndim * sizeof(shape_elem_type)); - - memcpy_event.wait(); - - sycl::range<2> gws(size_, kth + 1); - auto kernel_parallel_for_func = [=](sycl::id<2> global_id) { - size_t j = global_id[0]; - size_t k = global_id[1]; - - _DataType val = arr2[j * shape[ndim - 1] + k]; - - for (size_t i = 0; i < static_cast(shape[ndim - 1]); ++i) { - if (result[j * shape[ndim - 1] + i] == val) { - _DataType change_val1 = result[j * shape[ndim - 1] + i]; - _DataType change_val2 = result[j * shape[ndim - 1] + k]; - result[j * shape[ndim - 1] + k] = change_val1; - result[j * shape[ndim - 1] + i] = change_val2; - } - } - }; - - auto kernel_func = [&](sycl::handler &cgh) { - cgh.depends_on({memcpy_event}); - cgh.parallel_for>( - gws, kernel_parallel_for_func); - }; - - auto event = q.submit(kernel_func); - - event.wait(); - - delete[] matrix; - sycl::free(shape, q); - return event_ref; } diff --git a/dpnp/dpnp_array.py b/dpnp/dpnp_array.py index f37a3a2b3be3..bb864d4444a9 100644 --- a/dpnp/dpnp_array.py +++ b/dpnp/dpnp_array.py @@ -1459,21 +1459,34 @@ def nonzero(self): def partition(self, /, kth, axis=-1, kind="introselect", order=None): """ - Return a partitioned copy of an array. + Partially sorts the elements in the array in such a way that the value + of the element in k-th position is in the position it would be in a + sorted array. In the output array, all elements smaller than the k-th + element are located to the left of this element and all equal or + greater are located to its right. The ordering of the elements in the + two partitions on the either side of the k-th element in the output + array is undefined. - Rearranges the elements in the array in such a way that the value of - the element in `kth` position is in the position it would be in - a sorted array. + Refer to `dpnp.partition` for full documentation. - All elements smaller than the `kth` element are moved before this - element and all equal or greater are moved behind it. The ordering - of the elements in the two partitions is undefined. + kth : {int, sequence of ints} + Element index to partition by. The kth element value will be in its + final sorted position and all smaller elements will be moved before + it and all equal or greater elements behind it. + The order of all elements in the partitions is undefined. If + provided with a sequence of kth it will partition all elements + indexed by kth of them into their sorted position at once. + axis : int, optional + Axis along which to sort. The default is ``-1``, which means sort + along the the last axis. - Refer to `dpnp.partition` for full documentation. + Default: ``-1``. See Also -------- :obj:`dpnp.partition` : Return a partitioned copy of an array. + :obj:`dpnp.argpartition` : Indirect partition. + :obj:`dpnp.sort` : Full sort. Examples -------- @@ -1481,13 +1494,19 @@ def partition(self, /, kth, axis=-1, kind="introselect", order=None): >>> a = np.array([3, 4, 2, 1]) >>> a.partition(3) >>> a + array([1, 2, 3, 4]) # may vary + + >>> a.partition((1, 3)) + >>> a array([1, 2, 3, 4]) """ - self._array_obj = dpnp.partition( - self, kth, axis=axis, kind=kind, order=order - ).get_array() + if axis is None: + raise TypeError( + "'NoneType' object cannot be interpreted as an integer" + ) + self[...] = dpnp.partition(self, kth, axis=axis, kind=kind, order=order) def prod( self, diff --git a/dpnp/dpnp_iface_sorting.py b/dpnp/dpnp_iface_sorting.py index db33a88c7488..9c5097a5f3e3 100644 --- a/dpnp/dpnp_iface_sorting.py +++ b/dpnp/dpnp_iface_sorting.py @@ -39,8 +39,9 @@ """ +from collections.abc import Sequence + import dpctl.tensor as dpt -import numpy from dpctl.tensor._numpy_helper import normalize_axis_index import dpnp @@ -51,7 +52,6 @@ ) from .dpnp_array import dpnp_array from .dpnp_utils import ( - call_origin, map_dtype_to_device, ) @@ -147,7 +147,7 @@ def argsort( Limitations ----------- - Parameters `order` is only supported with its default value. + Parameter `order` is only supported with its default value. Otherwise ``NotImplementedError`` exception will be raised. Sorting algorithms ``"quicksort"`` and ``"heapsort"`` are not supported. @@ -201,44 +201,128 @@ def argsort( ) -def partition(x1, kth, axis=-1, kind="introselect", order=None): +def partition(a, kth, axis=-1, kind="introselect", order=None): """ Return a partitioned copy of an array. For full documentation refer to :obj:`numpy.partition`. + Parameters + ---------- + a : {dpnp.ndarray, usm_ndarray} + Array to be sorted. + kth : {int, sequence of ints} + Element index to partition by. The k-th value of the element will be in + its final sorted position and all smaller elements will be moved before + it and all equal or greater elements behind it. The order of all + elements in the partitions is undefined. If provided with a sequence of + k-th it will partition all elements indexed by k-th of them into their + sorted position at once. + 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``. + + Returns + ------- + out : dpnp.ndarray + Array of the same type and shape as `a`. + Limitations ----------- - Input array is supported as :obj:`dpnp.ndarray`. - Input `kth` is supported as :obj:`int`. - Parameters `axis`, `kind` and `order` are supported only with default - values. + Parameters `kind` and `order` are only supported with its default value. + Otherwise ``NotImplementedError`` exception will be raised. + + See Also + -------- + :obj:`dpnp.ndarray.partition` : Equivalent method. + :obj:`dpnp.argpartition` : Indirect partition. + :obj:`dpnp.sort` : Full sorting. + + Examples + -------- + >>> import dpnp as np + >>> a = np.array([7, 1, 7, 7, 1, 5, 7, 2, 3, 2, 6, 2, 3, 0]) + >>> p = np.partition(a, 4) + >>> p + array([0, 1, 1, 2, 2, 2, 3, 3, 5, 7, 7, 7, 7, 6]) # may vary + + ``p[4]`` is 2; all elements in ``p[:4]`` are less than or equal to + ``p[4]``, and all elements in ``p[5:]`` are greater than or equal to + ``p[4]``. The partition is:: + + [0, 1, 1, 2], [2], [2, 3, 3, 5, 7, 7, 7, 7, 6] + + The next example shows the use of multiple values passed to `kth`. + + >>> p2 = np.partition(a, (4, 8)) + >>> p2 + array([0, 1, 1, 2, 2, 2, 3, 3, 5, 6, 7, 7, 7, 7]) + + ``p2[4]`` is 2 and ``p2[8]`` is 5. All elements in ``p2[:4]`` are less + than or equal to ``p2[4]``, all elements in ``p2[5:8]`` are greater than or + equal to ``p2[4]`` and less than or equal to ``p2[8]``, and all elements in + ``p2[9:]`` are greater than or equal to ``p2[8]``. The partition is:: + + [0, 1, 1, 2], [2], [2, 3, 3], [5], [6, 7, 7, 7, 7] """ - x1_desc = dpnp.get_dpnp_descriptor(x1, copy_when_nondefault_queue=False) - if x1_desc: - if dpnp.is_cuda_backend(x1_desc.get_array()): # pragma: no cover - raise NotImplementedError( - "Running on CUDA is currently not supported" - ) + dpnp.check_supported_arrays_type(a) - if not isinstance(kth, int): - pass - elif x1_desc.ndim == 0: - pass - elif kth >= x1_desc.shape[x1_desc.ndim - 1] or x1_desc.ndim + kth < 0: - pass - elif axis != -1: - pass - elif kind != "introselect": - pass - elif order is not None: - pass - else: - return dpnp_partition(x1_desc, kth, axis, kind, order).get_pyobj() + if kind != "introselect": + raise NotImplementedError( + "`kind` keyword argument is only supported with its default value." + ) + if order is not None: + raise NotImplementedError( + "`order` keyword argument is only supported with its default value." + ) - return call_origin(numpy.partition, x1, kth, axis, kind, order) + if axis is None: + a = dpnp.ravel(a) + axis = -1 + + nd = a.ndim + axis = normalize_axis_index(axis, nd) + length = a.shape[axis] + + if isinstance(kth, int): + kth = (kth,) + elif not isinstance(kth, Sequence): + raise TypeError( + f"kth must be int or sequence of ints, but got {type(kth)}" + ) + elif not all(isinstance(k, int) for k in kth): + raise TypeError("kth is a sequence, but not all elements are integers") + + nkth = len(kth) + if nkth == 0 or a.size == 0: + return dpnp.copy(a) + + # validate kth + kth = list(kth) + for i in range(nkth): + if kth[i] < 0: + kth[i] += length + + if not 0 <= kth[i] < length: + raise ValueError(f"kth(={kth[i]}) out of bounds {length}") + + dt = a.dtype + if ( + nd > 1 + or nkth > 1 + or dpnp.issubdtype(dt, dpnp.unsignedinteger) + or dt in (dpnp.int8, dpnp.int16) + or dpnp.is_cuda_backend(a.get_array()) + ): + # sort is a faster path in case of ndim > 1 + return dpnp.sort(a, axis=axis) + + desc = dpnp.get_dpnp_descriptor(a, copy_when_nondefault_queue=False) + return dpnp_partition(desc, kth[0], axis, kind, order).get_pyobj() def sort(a, axis=-1, kind=None, order=None, *, descending=False, stable=None): diff --git a/dpnp/tests/skipped_tests_cuda.tbl b/dpnp/tests/skipped_tests_cuda.tbl index f035fba3302f..e8415a1ae410 100644 --- a/dpnp/tests/skipped_tests_cuda.tbl +++ b/dpnp/tests/skipped_tests_cuda.tbl @@ -660,144 +660,3 @@ tests/third_party/cupy/random_tests/test_sample.py::TestRandomIntegers2::test_bo tests/third_party/cupy/random_tests/test_sample.py::TestRandomIntegers2::test_bound_2 tests/third_party/cupy/random_tests/test_sample.py::TestRandomIntegers2::test_goodness_of_fit tests/third_party/cupy/random_tests/test_sample.py::TestRandomIntegers2::test_goodness_of_fit_2 - -# partition -tests/test_sort.py::test_partition[[3, 4, 2, 1]-bool-0] -tests/test_sort.py::test_partition[[3, 4, 2, 1]-bool-1] -tests/test_sort.py::test_partition[[3, 4, 2, 1]-int32-0] -tests/test_sort.py::test_partition[[3, 4, 2, 1]-int32-1] -tests/test_sort.py::test_partition[[3, 4, 2, 1]-int64-0] -tests/test_sort.py::test_partition[[3, 4, 2, 1]-int64-1] -tests/test_sort.py::test_partition[[3, 4, 2, 1]-float32-0] -tests/test_sort.py::test_partition[[3, 4, 2, 1]-float32-1] -tests/test_sort.py::test_partition[[3, 4, 2, 1]-float64-0] -tests/test_sort.py::test_partition[[3, 4, 2, 1]-float64-1] -tests/test_sort.py::test_partition[[3, 4, 2, 1]-complex64-0] -tests/test_sort.py::test_partition[[3, 4, 2, 1]-complex64-1] -tests/test_sort.py::test_partition[[3, 4, 2, 1]-complex128-0] -tests/test_sort.py::test_partition[[3, 4, 2, 1]-complex128-1] -tests/test_sort.py::test_partition[[[1, 0], [3, 0]]-bool-0] -tests/test_sort.py::test_partition[[[1, 0], [3, 0]]-bool-1] -tests/test_sort.py::test_partition[[[1, 0], [3, 0]]-int32-0] -tests/test_sort.py::test_partition[[[1, 0], [3, 0]]-int32-1] -tests/test_sort.py::test_partition[[[1, 0], [3, 0]]-int64-0] -tests/test_sort.py::test_partition[[[1, 0], [3, 0]]-int64-1] -tests/test_sort.py::test_partition[[[1, 0], [3, 0]]-float32-0] -tests/test_sort.py::test_partition[[[1, 0], [3, 0]]-float32-1] -tests/test_sort.py::test_partition[[[1, 0], [3, 0]]-float64-0] -tests/test_sort.py::test_partition[[[1, 0], [3, 0]]-float64-1] -tests/test_sort.py::test_partition[[[1, 0], [3, 0]]-complex64-0] -tests/test_sort.py::test_partition[[[1, 0], [3, 0]]-complex64-1] -tests/test_sort.py::test_partition[[[1, 0], [3, 0]]-complex128-0] -tests/test_sort.py::test_partition[[[1, 0], [3, 0]]-complex128-1] -tests/test_sort.py::test_partition[[[3, 2], [1, 6]]-bool-0] -tests/test_sort.py::test_partition[[[3, 2], [1, 6]]-bool-1] -tests/test_sort.py::test_partition[[[3, 2], [1, 6]]-int32-0] -tests/test_sort.py::test_partition[[[3, 2], [1, 6]]-int32-1] -tests/test_sort.py::test_partition[[[3, 2], [1, 6]]-int64-0] -tests/test_sort.py::test_partition[[[3, 2], [1, 6]]-int64-1] -tests/test_sort.py::test_partition[[[3, 2], [1, 6]]-float32-0] -tests/test_sort.py::test_partition[[[3, 2], [1, 6]]-float32-1] -tests/test_sort.py::test_partition[[[3, 2], [1, 6]]-float64-0] -tests/test_sort.py::test_partition[[[3, 2], [1, 6]]-float64-1] -tests/test_sort.py::test_partition[[[3, 2], [1, 6]]-complex64-0] -tests/test_sort.py::test_partition[[[3, 2], [1, 6]]-complex64-1] -tests/test_sort.py::test_partition[[[3, 2], [1, 6]]-complex128-0] -tests/test_sort.py::test_partition[[[3, 2], [1, 6]]-complex128-1] -tests/test_sort.py::test_partition[[[4, 2, 3], [3, 4, 1]]-bool-0] -tests/test_sort.py::test_partition[[[4, 2, 3], [3, 4, 1]]-bool-1] -tests/test_sort.py::test_partition[[[4, 2, 3], [3, 4, 1]]-int32-0] -tests/test_sort.py::test_partition[[[4, 2, 3], [3, 4, 1]]-int32-1] -tests/test_sort.py::test_partition[[[4, 2, 3], [3, 4, 1]]-int64-0] -tests/test_sort.py::test_partition[[[4, 2, 3], [3, 4, 1]]-int64-1] -tests/test_sort.py::test_partition[[[4, 2, 3], [3, 4, 1]]-float32-0] -tests/test_sort.py::test_partition[[[4, 2, 3], [3, 4, 1]]-float32-1] -tests/test_sort.py::test_partition[[[4, 2, 3], [3, 4, 1]]-float64-0] -tests/test_sort.py::test_partition[[[4, 2, 3], [3, 4, 1]]-float64-1] -tests/test_sort.py::test_partition[[[4, 2, 3], [3, 4, 1]]-complex64-0] -tests/test_sort.py::test_partition[[[4, 2, 3], [3, 4, 1]]-complex64-1] -tests/test_sort.py::test_partition[[[4, 2, 3], [3, 4, 1]]-complex128-0] -tests/test_sort.py::test_partition[[[4, 2, 3], [3, 4, 1]]-complex128-1] -tests/test_sort.py::test_partition[[[[1, -3], [3, 0]], [[5, 2], [0, 1]], [[1, 0], [0, 1]]]-bool-0] -tests/test_sort.py::test_partition[[[[1, -3], [3, 0]], [[5, 2], [0, 1]], [[1, 0], [0, 1]]]-bool-1] -tests/test_sort.py::test_partition[[[[1, -3], [3, 0]], [[5, 2], [0, 1]], [[1, 0], [0, 1]]]-int32-0] -tests/test_sort.py::test_partition[[[[1, -3], [3, 0]], [[5, 2], [0, 1]], [[1, 0], [0, 1]]]-int32-1] -tests/test_sort.py::test_partition[[[[1, -3], [3, 0]], [[5, 2], [0, 1]], [[1, 0], [0, 1]]]-int64-0] -tests/test_sort.py::test_partition[[[[1, -3], [3, 0]], [[5, 2], [0, 1]], [[1, 0], [0, 1]]]-int64-1] -tests/test_sort.py::test_partition[[[[1, -3], [3, 0]], [[5, 2], [0, 1]], [[1, 0], [0, 1]]]-float32-0] -tests/test_sort.py::test_partition[[[[1, -3], [3, 0]], [[5, 2], [0, 1]], [[1, 0], [0, 1]]]-float32-1] -tests/test_sort.py::test_partition[[[[1, -3], [3, 0]], [[5, 2], [0, 1]], [[1, 0], [0, 1]]]-float64-0] -tests/test_sort.py::test_partition[[[[1, -3], [3, 0]], [[5, 2], [0, 1]], [[1, 0], [0, 1]]]-float64-1] -tests/test_sort.py::test_partition[[[[1, -3], [3, 0]], [[5, 2], [0, 1]], [[1, 0], [0, 1]]]-complex64-0] -tests/test_sort.py::test_partition[[[[1, -3], [3, 0]], [[5, 2], [0, 1]], [[1, 0], [0, 1]]]-complex64-1] -tests/test_sort.py::test_partition[[[[1, -3], [3, 0]], [[5, 2], [0, 1]], [[1, 0], [0, 1]]]-complex128-0] -tests/test_sort.py::test_partition[[[[1, -3], [3, 0]], [[5, 2], [0, 1]], [[1, 0], [0, 1]]]-complex128-1] -tests/test_sort.py::test_partition[[[[[8, 2], [3, 0]], [[5, 2], [0, 1]]], [[[1, 3], [3, 1]], [[5, 2], [0, 1]]]]-bool-0] -tests/test_sort.py::test_partition[[[[[8, 2], [3, 0]], [[5, 2], [0, 1]]], [[[1, 3], [3, 1]], [[5, 2], [0, 1]]]]-bool-1] -tests/test_sort.py::test_partition[[[[[8, 2], [3, 0]], [[5, 2], [0, 1]]], [[[1, 3], [3, 1]], [[5, 2], [0, 1]]]]-int32-0] -tests/test_sort.py::test_partition[[[[[8, 2], [3, 0]], [[5, 2], [0, 1]]], [[[1, 3], [3, 1]], [[5, 2], [0, 1]]]]-int32-1] -tests/test_sort.py::test_partition[[[[[8, 2], [3, 0]], [[5, 2], [0, 1]]], [[[1, 3], [3, 1]], [[5, 2], [0, 1]]]]-int64-0] -tests/test_sort.py::test_partition[[[[[8, 2], [3, 0]], [[5, 2], [0, 1]]], [[[1, 3], [3, 1]], [[5, 2], [0, 1]]]]-int64-1] -tests/test_sort.py::test_partition[[[[[8, 2], [3, 0]], [[5, 2], [0, 1]]], [[[1, 3], [3, 1]], [[5, 2], [0, 1]]]]-float32-0] -tests/test_sort.py::test_partition[[[[[8, 2], [3, 0]], [[5, 2], [0, 1]]], [[[1, 3], [3, 1]], [[5, 2], [0, 1]]]]-float32-1] -tests/test_sort.py::test_partition[[[[[8, 2], [3, 0]], [[5, 2], [0, 1]]], [[[1, 3], [3, 1]], [[5, 2], [0, 1]]]]-float64-0] -tests/test_sort.py::test_partition[[[[[8, 2], [3, 0]], [[5, 2], [0, 1]]], [[[1, 3], [3, 1]], [[5, 2], [0, 1]]]]-float64-1] -tests/test_sort.py::test_partition[[[[[8, 2], [3, 0]], [[5, 2], [0, 1]]], [[[1, 3], [3, 1]], [[5, 2], [0, 1]]]]-complex64-0] -tests/test_sort.py::test_partition[[[[[8, 2], [3, 0]], [[5, 2], [0, 1]]], [[[1, 3], [3, 1]], [[5, 2], [0, 1]]]]-complex64-1] -tests/test_sort.py::test_partition[[[[[8, 2], [3, 0]], [[5, 2], [0, 1]]], [[[1, 3], [3, 1]], [[5, 2], [0, 1]]]]-complex128-0] -tests/test_sort.py::test_partition[[[[[8, 2], [3, 0]], [[5, 2], [0, 1]]], [[[1, 3], [3, 1]], [[5, 2], [0, 1]]]]-complex128-1] - -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_0_{external=False, length=10}::test_partition_axis -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_0_{external=False, length=10}::test_partition_invalid_axis1 -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_0_{external=False, length=10}::test_partition_invalid_axis2 -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_0_{external=False, length=10}::test_partition_invalid_kth -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_0_{external=False, length=10}::test_partition_invalid_negative_axis1 -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_0_{external=False, length=10}::test_partition_invalid_negative_axis2 -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_0_{external=False, length=10}::test_partition_invalid_negative_kth -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_0_{external=False, length=10}::test_partition_negative_axis -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_0_{external=False, length=10}::test_partition_negative_kth -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_0_{external=False, length=10}::test_partition_non_contiguous -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_0_{external=False, length=10}::test_partition_one_dim -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_0_{external=False, length=10}::test_partition_sequence_kth -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_0_{external=False, length=10}::test_partition_zero_dim -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_1_{external=False, length=20000}::test_partition_axis -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_1_{external=False, length=20000}::test_partition_invalid_axis1 -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_1_{external=False, length=20000}::test_partition_invalid_axis2 -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_1_{external=False, length=20000}::test_partition_invalid_kth -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_1_{external=False, length=20000}::test_partition_invalid_negative_axis1 -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_1_{external=False, length=20000}::test_partition_invalid_negative_axis2 -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_1_{external=False, length=20000}::test_partition_invalid_negative_kth -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_1_{external=False, length=20000}::test_partition_negative_axis -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_1_{external=False, length=20000}::test_partition_negative_kth -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_1_{external=False, length=20000}::test_partition_non_contiguous -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_1_{external=False, length=20000}::test_partition_one_dim -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_1_{external=False, length=20000}::test_partition_sequence_kth -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_1_{external=False, length=20000}::test_partition_zero_dim -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_2_{external=True, length=10}::test_partition_axis -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_2_{external=True, length=10}::test_partition_invalid_axis1 -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_2_{external=True, length=10}::test_partition_invalid_axis2 -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_2_{external=True, length=10}::test_partition_invalid_kth -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_2_{external=True, length=10}::test_partition_invalid_negative_axis1 -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_2_{external=True, length=10}::test_partition_invalid_negative_axis2 -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_2_{external=True, length=10}::test_partition_invalid_negative_kth -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_2_{external=True, length=10}::test_partition_negative_axis -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_2_{external=True, length=10}::test_partition_negative_kth -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_2_{external=True, length=10}::test_partition_non_contiguous -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_2_{external=True, length=10}::test_partition_none_axis -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_2_{external=True, length=10}::test_partition_one_dim -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_2_{external=True, length=10}::test_partition_sequence_kth -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_2_{external=True, length=10}::test_partition_zero_dim -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_3_{external=True, length=20000}::test_partition_axis -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_3_{external=True, length=20000}::test_partition_invalid_axis1 -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_3_{external=True, length=20000}::test_partition_invalid_axis2 -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_3_{external=True, length=20000}::test_partition_invalid_kth -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_3_{external=True, length=20000}::test_partition_invalid_negative_axis1 -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_3_{external=True, length=20000}::test_partition_invalid_negative_axis2 -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_3_{external=True, length=20000}::test_partition_invalid_negative_kth -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_3_{external=True, length=20000}::test_partition_negative_axis -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_3_{external=True, length=20000}::test_partition_negative_kth -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_3_{external=True, length=20000}::test_partition_non_contiguous -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_3_{external=True, length=20000}::test_partition_none_axis -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_3_{external=True, length=20000}::test_partition_one_dim -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_3_{external=True, length=20000}::test_partition_sequence_kth -tests/third_party/cupy/sorting_tests/test_sort.py::TestPartition_param_3_{external=True, length=20000}::test_partition_zero_dim diff --git a/dpnp/tests/test_sort.py b/dpnp/tests/test_sort.py index 31959600a0d7..5e883c575f85 100644 --- a/dpnp/tests/test_sort.py +++ b/dpnp/tests/test_sort.py @@ -11,6 +11,7 @@ get_all_dtypes, get_complex_dtypes, get_float_dtypes, + get_integer_dtypes, ) from .third_party.cupy import testing @@ -275,6 +276,232 @@ def test_v_scalar(self): assert_equal(result, expected) +class TestPartition: + @pytest.mark.parametrize("data", [[2, 1], [1, 2], [1, 1]]) + @pytest.mark.parametrize("kth", [0, 1]) + def test_1d_2size(self, data, kth): + a = numpy.array(data) + ia = dpnp.array(a) + + result = dpnp.partition(ia, kth) + expected = numpy.partition(a, kth) + assert_array_equal(result, expected) + + @pytest.mark.parametrize( + "data", + [ + [3, 2, 1], + [1, 2, 3], + [2, 1, 3], + [2, 3, 1], + [1, 1, 1], + [1, 2, 2], + [2, 2, 1], + [1, 2, 1], + ], + ) + @pytest.mark.parametrize("kth", [0, 1, 2]) + @pytest.mark.parametrize("dt", get_all_dtypes(no_none=True)) + def test_1d_3size(self, data, kth, dt): + a = dpnp.array(data, dtype=dt) + p = dpnp.partition(a, kth) + + assert (p[..., 0:kth] <= p[..., kth : kth + 1]).all() + assert (p[..., kth : kth + 1] <= p[..., kth + 1 :]).all() + + @pytest.mark.parametrize("kth", [6, 16, -6, 41, -16, 31]) + def test_1d_reversed(self, kth): + a = dpnp.arange(47)[::-1] + p = dpnp.partition(a, kth) + + assert (p[..., 0:kth] <= p[..., kth : kth + 1]).all() + assert (p[..., kth : kth + 1] <= p[..., kth + 1 :]).all() + + @pytest.mark.parametrize("val", [4, dpnp.nan]) + def test_1d_ones(self, val): + a = numpy.ones(10) + a[1] = val + ia = dpnp.array(a) + + result = dpnp.partition(ia, (2, -1)) + expected = numpy.partition(a, (2, -1)) + assert_array_equal(result, expected) + + @pytest.mark.parametrize("kth", [0, 3, 19, 20]) + def test_1d_equal_elements(self, kth): + a = dpnp.array( + [ + 0, + 1, + 2, + 3, + 4, + 5, + 7, + 7, + 7, + 7, + 7, + 7, + 7, + 7, + 7, + 7, + 7, + 7, + 7, + 7, + 7, + 7, + 9, + ] + ) + p = dpnp.partition(a, kth) + + assert (p[..., 0:kth] <= p[..., kth : kth + 1]).all() + assert (p[..., kth : kth + 1] <= p[..., kth + 1 :]).all() + + @pytest.mark.parametrize("kth", [(0, 3), (-3, -1)]) + def test_kth_iterative(self, kth): + a = numpy.array([3, 4, 2, 1]) + ia = dpnp.array(a) + + result = dpnp.partition(ia, kth) + expected = numpy.partition(a, kth) + assert_array_equal(result, expected) + + @pytest.mark.parametrize("dt", get_integer_dtypes()) + def test_max_min_int_values(self, dt): + N = 512 + rnd = numpy.random.RandomState(1100710816) + + # random data with min and max values + minv = numpy.iinfo(dt).min + maxv = numpy.iinfo(dt).max + a = rnd.randint(low=minv, high=maxv, size=N, dtype=dt) + i, j = rnd.choice(N, 2, replace=False) + a[i] = minv + a[j] = maxv + k = int(rnd.choice(N, 1)[0]) + + ia = dpnp.array(a) + p = dpnp.partition(ia, k) + assert (p[0:k] <= p[k : k + 1]).all() + assert (p[k : k + 1] <= p[k + 1 :]).all() + + # random data with max value at the end of array + a = rnd.randint(low=minv, high=maxv, size=N, dtype=dt) + a[N - 1] = maxv + + ia = dpnp.array(a) + p = dpnp.partition(ia, k) + assert (p[0:k] <= p[k : k + 1]).all() + assert (p[k : k + 1] <= p[k + 1 :]).all() + + @pytest.mark.parametrize("dt", get_float_dtypes()) + def test_float_values(self, dt): + N = 512 + rnd = numpy.random.RandomState(1100710816) + a = -0.5 + rnd.random(N).astype(dt) + k = int(rnd.choice(N, 1)[0]) + + ia = dpnp.array(a) + p = dpnp.partition(ia, k) + assert (p[0:k] <= p[k : k + 1]).all() + assert (p[k : k + 1] <= p[k + 1 :]).all() + + @pytest.mark.parametrize("axis", [0, -1, None]) + def test_axis_1d(self, axis): + a = numpy.array([2, 1]) + ia = dpnp.array(a) + + result = dpnp.partition(ia, 1, axis=axis) + expected = numpy.partition(a, 1, axis=axis) + assert_array_equal(result, expected) + + @pytest.mark.parametrize("kth, axis", [(1, 0), (4, 1)]) + def test_axis_2d(self, kth, axis): + a = generate_random_numpy_array((2, 5)) + + ia = dpnp.array(a) + ia.partition(kth, axis=axis) + p = dpnp.rollaxis(ia, axis, ia.ndim) + assert (p[..., 0:kth] <= p[..., kth : kth + 1]).all() + assert (p[..., kth : kth + 1] <= p[..., kth + 1 :]).all() + + ia = dpnp.array(a) + p = dpnp.partition(ia, kth, axis=axis) + p = dpnp.rollaxis(p, axis, ia.ndim) + assert (p[..., 0:kth] <= p[..., kth : kth + 1]).all() + assert (p[..., kth : kth + 1] <= p[..., kth + 1 :]).all() + + @pytest.mark.parametrize("kth", [1, 9]) + def test_axis_2d_none(self, kth): + a = generate_random_numpy_array((2, 5)) + ia = dpnp.array(a) + + p = dpnp.partition(ia, kth, axis=None) + assert (p[..., 0:kth] <= p[..., kth : kth + 1]).all() + assert (p[..., kth : kth + 1] <= p[..., kth + 1 :]).all() + + @pytest.mark.parametrize("axis", list(range(-4, 4)) + [None]) + def test_empty_array(self, axis): + a = numpy.empty((3, 2, 1, 0)) + ia = dpnp.array(a) + kth = 0 + + result = dpnp.partition(ia, kth, axis=axis) + expected = numpy.partition(a, kth, axis=axis) + assert_equal(result, expected) + + def test_empty_partition(self): + a = numpy.array([0, 2, 4, 6, 8, 10]) + ia = dpnp.array(a) + + ia.partition([]) + assert_array_equal(ia, a) + + @pytest.mark.parametrize("xp", [dpnp, numpy]) + def test_kth_errors(self, xp): + a = xp.arange(10) + assert_raises(ValueError, a.partition, 10) + assert_raises(ValueError, a.partition, -11) + assert_raises(TypeError, a.partition, 9.0) + assert_raises(TypeError, a.partition, [1, 7.0]) + + @pytest.mark.parametrize("xp", [dpnp, numpy]) + def test_kth_axis_errors(self, xp): + a = xp.array([2, 1]) + assert_raises(ValueError, a.partition, 2) + assert_raises(AxisError, a.partition, 3, axis=1) + assert_raises(ValueError, xp.partition, a, 2) + assert_raises(AxisError, xp.partition, a, 2, axis=1) + + a = xp.arange(10).reshape((2, 5)) + assert_raises(ValueError, a.partition, 2, axis=0) + assert_raises(ValueError, a.partition, 11, axis=1) + assert_raises(TypeError, a.partition, 2, axis=None) + assert_raises(ValueError, xp.partition, a, 9, axis=1) + assert_raises(ValueError, xp.partition, a, 11, axis=None) + + @pytest.mark.parametrize("xp", [dpnp, numpy]) + def test_kth_iterative_error(self, xp): + a = xp.arange(17) + kth = (0, 1, 2, 429, 231) + assert_raises(ValueError, a.partition, kth) + + a = xp.arange(10).reshape((2, 5)) + assert_raises(ValueError, a.partition, kth, axis=0) + assert_raises(ValueError, a.partition, kth, axis=1) + assert_raises(ValueError, xp.partition, a, kth, axis=1) + assert_raises(ValueError, xp.partition, a, kth, axis=None) + + def test_not_implemented_kwargs(self): + a = dpnp.arange(10) + assert_raises(NotImplementedError, a.partition, 2, kind="nonsense") + assert_raises(NotImplementedError, a.partition, 2, order=[]) + + class TestSort: @pytest.mark.parametrize("kind", [None, "stable", "mergesort", "radixsort"]) @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True)) @@ -407,40 +634,3 @@ def test_complex(self, dtype): result = dpnp.sort_complex(ia) expected = numpy.sort_complex(a) assert_equal(result, expected) - - -@pytest.mark.parametrize("kth", [0, 1]) -@pytest.mark.parametrize( - "dtype", - get_all_dtypes( - no_none=True, no_unsigned=True, xfail_dtypes=[dpnp.int8, dpnp.int16] - ), -) -@pytest.mark.parametrize( - "array", - [ - [3, 4, 2, 1], - [[1, 0], [3, 0]], - [[3, 2], [1, 6]], - [[4, 2, 3], [3, 4, 1]], - [[[1, -3], [3, 0]], [[5, 2], [0, 1]], [[1, 0], [0, 1]]], - [ - [[[8, 2], [3, 0]], [[5, 2], [0, 1]]], - [[[1, 3], [3, 1]], [[5, 2], [0, 1]]], - ], - ], - ids=[ - "[3, 4, 2, 1]", - "[[1, 0], [3, 0]]", - "[[3, 2], [1, 6]]", - "[[4, 2, 3], [3, 4, 1]]", - "[[[1, -3], [3, 0]], [[5, 2], [0, 1]], [[1, 0], [0, 1]]]", - "[[[[8, 2], [3, 0]], [[5, 2], [0, 1]]], [[[1, 3], [3, 1]], [[5, 2], [0, 1]]]]", - ], -) -def test_partition(array, dtype, kth): - a = dpnp.array(array, dtype) - p = dpnp.partition(a, kth) - - assert (p[..., 0:kth] <= p[..., kth : kth + 1]).all() - assert (p[..., kth : kth + 1] <= p[..., kth + 1 :]).all() diff --git a/dpnp/tests/third_party/cupy/sorting_tests/test_sort.py b/dpnp/tests/third_party/cupy/sorting_tests/test_sort.py index 7e791f6a7c0e..7e0eade13254 100644 --- a/dpnp/tests/third_party/cupy/sorting_tests/test_sort.py +++ b/dpnp/tests/third_party/cupy/sorting_tests/test_sort.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import unittest import numpy @@ -455,7 +457,6 @@ def test_sort_complex_nan(self, xp, dtype): } ) ) -@pytest.mark.usefixtures("allow_fall_back_on_numpy") class TestPartition(unittest.TestCase): def partition(self, a, kth, axis=-1): @@ -478,9 +479,6 @@ def test_partition_zero_dim(self): @testing.for_all_dtypes() @testing.numpy_cupy_equal() def test_partition_one_dim(self, xp, dtype): - flag = xp.issubdtype(dtype, xp.unsignedinteger) - if flag or dtype in [xp.int8, xp.int16]: - pytest.skip("dpnp.partition() does not support new integer dtypes.") a = testing.shaped_random((self.length,), xp, dtype) kth = 2 x = self.partition(a, kth) @@ -488,7 +486,6 @@ def test_partition_one_dim(self, xp, dtype): assert xp.all(x[kth : kth + 1] <= x[kth + 1 :]) return x[kth] - @pytest.mark.skip("multidimensional case doesn't work properly") @testing.for_all_dtypes() @testing.numpy_cupy_array_equal() def test_partition_multi_dim(self, xp, dtype): @@ -505,6 +502,12 @@ def test_partition_multi_dim(self, xp, dtype): def test_partition_non_contiguous(self, xp): a = testing.shaped_random((self.length,), xp)[::-1] kth = 2 + # if not self.external: + # if xp is cupy: + # with self.assertRaises(NotImplementedError): + # return self.partition(a, kth) + # return 0 # dummy + # else: x = self.partition(a, kth) assert xp.all(x[0:kth] <= x[kth : kth + 1]) assert xp.all(x[kth : kth + 1] <= x[kth + 1 :]) @@ -607,7 +610,7 @@ def test_partition_invalid_negative_axis2(self): } ) ) -@pytest.mark.skip("not fully supported yet") +@pytest.mark.skip("not supported yet") class TestArgpartition(unittest.TestCase): def argpartition(self, a, kth, axis=-1):