Implement kernels for median and paritition functions

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 * 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)
+* Added `dpnp.argpartition` and `dpnp.ndarray.argpartition` to return the indices that would partition an array [#3138](https://github.com/IntelPython/dpnp/pull/3138)
 
 ### Changed
 
@@ -18,6 +19,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Merged the `ufunc` and `indexing` extensions into `dpnp.libtensor._tensor_elementwise_impl` and `dpnp.libtensor._tensor_impl`, while `statistics` and `window` extensions became part of libtensor [#3131](https://github.com/IntelPython/dpnp/pull/3131)
 * Moved `blas`, `fft`, `lapack`, `sparse` and `vm` extensions into `dpnp/libtensor/extensions` [#3131](https://github.com/IntelPython/dpnp/pull/3131)
 * Moved the libtensor headers shipped with dpnp from `dpnp/tensor/libtensor/include` to `dpnp/libtensor/include`, and changed `dpnp.get_include` to return that directory [#3131](https://github.com/IntelPython/dpnp/pull/3131)
+* Changed `dpnp.partition`, `dpnp.median` and `dpnp.nanmedian` to use a radix select instead of a full sort for all but complex data types [#3138](https://github.com/IntelPython/dpnp/pull/3138)
 
 ### Deprecated
 
@@ -29,12 +31,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Removed
 
 * Removed the unused `FourZeroOffsets_Indexer`, `FixedDimStridedIndexer` and `ThreeOffsets_FixedDimStridedIndexer` from the libtensor header `utils/offset_utils.hpp` [#3131](https://github.com/IntelPython/dpnp/pull/3131)
+* Removed the legacy implementation of `dpnp.partition` from the backend [#3138](https://github.com/IntelPython/dpnp/pull/3138)
 
 ### Fixed
 
 * Fixed `flags[key]` of `dpnp.ndarray` and `dpnp.tensor.usm_ndarray` to raise `KeyError` for an unknown key [#3141](https://github.com/IntelPython/dpnp/pull/3141)
 * Released the GIL around the oneMKL VM calls in the VM extension, which deadlocked on an in-order queue of a CPU device [#3156](https://github.com/IntelPython/dpnp/pull/3156)
 * Fixed `dpnp.place` and `dpnp.tensor.place` to order the write after the cast of `vals` to the array's data type, which could put wrong values into the array on an out-of-order queue [#3164](https://github.com/IntelPython/dpnp/pull/3164)
+* Fixed `dpnp.median` and `dpnp.nanmedian` raising `ValueError` for an `out` array of the result shape with `keepdims=True` [#3138](https://github.com/IntelPython/dpnp/pull/3138)
 
 ### Security
 

dpnp/__init__.py

@@ -511,6 +511,7 @@
     where,
 )
 from .dpnp_iface_sorting import (
+    argpartition,
     argsort,
     partition,
     sort,
@@ -988,6 +989,7 @@
 __all__ += [
     "argmax",
     "argmin",
+    "argpartition",
     "argwhere",
     "argsort",
     "count_nonzero",

dpnp/backend/CMakeLists.txt

@@ -31,7 +31,6 @@ set(DPNP_SRC
     kernels/dpnp_krnl_arraycreation.cpp
     kernels/dpnp_krnl_common.cpp
     kernels/dpnp_krnl_random.cpp
-    kernels/dpnp_krnl_sorting.cpp
     src/dpnp_iface_fptr.cpp
     src/memory_sycl.cpp
     src/queue_sycl.cpp

dpnp/backend/include/dpnp_iface.hpp

@@ -103,38 +103,6 @@ INP_DLLEXPORT void dpnp_memory_memcpy_c(DPCTLSyclQueueRef q_ref,
 INP_DLLEXPORT void
     dpnp_memory_memcpy_c(void *dst, const void *src, size_t size_in_bytes);
 
-/**
- * @ingroup BACKEND_API
- * @brief Return a partitioned copy of an array.
- *
- * @param [in]  q_ref             Reference to SYCL queue.
- * @param [in]  array             Input array.
- * @param [in]  array2            Copy input array.
- * @param [out] result            Result array.
- * @param [in]  kth               Element index to partition by.
- * @param [in]  shape             Shape of input array.
- * @param [in]  ndim              Number of elements in shape.
- * @param [in]  dep_event_vec_ref Reference to vector of SYCL events.
- */
-template <typename _DataType>
-INP_DLLEXPORT DPCTLSyclEventRef
-    dpnp_partition_c(DPCTLSyclQueueRef q_ref,
-                     void *array,
-                     void *array2,
-                     void *result,
-                     const size_t kth,
-                     const shape_elem_type *shape,
-                     const size_t ndim,
-                     const DPCTLEventVectorRef dep_event_vec_ref);
-
-template <typename _DataType>
-INP_DLLEXPORT void dpnp_partition_c(void *array,
-                                    void *array2,
-                                    void *result,
-                                    const size_t kth,
-                                    const shape_elem_type *shape,
-                                    const size_t ndim);
-
 /**
  * @ingroup BACKEND_API
  * @brief implementation of creating filled with value array function

dpnp/backend/include/dpnp_iface_fptr.hpp

@@ -64,17 +64,14 @@ enum class DPNPFuncName : size_t
     DPNP_FN_NONE,    /**< Very first element of the enumeration */
     DPNP_FN_INITVAL, /**< Used in numpy ones, ones_like, zeros, zeros_like impls
                       */
-    DPNP_FN_INITVAL_EXT,   /**< Used in numpy ones, ones_like, zeros, zeros_like
-                              impls  */
-    DPNP_FN_ONES,          /**< Used in numpy.ones() impl */
-    DPNP_FN_ONES_LIKE,     /**< Used in numpy.ones_like() impl */
-    DPNP_FN_PARTITION,     /**< Used in numpy.partition() impl */
-    DPNP_FN_PARTITION_EXT, /**< Used in numpy.partition() impl, requires extra
-                              parameters */
-    DPNP_FN_RNG_BETA,      /**< Used in numpy.random.beta() impl  */
-    DPNP_FN_RNG_BETA_EXT,  /**< Used in numpy.random.beta() impl, requires extra
-                              parameters */
-    DPNP_FN_RNG_BINOMIAL,  /**< Used in numpy.random.binomial() impl  */
+    DPNP_FN_INITVAL_EXT,  /**< Used in numpy ones, ones_like, zeros, zeros_like
+                             impls  */
+    DPNP_FN_ONES,         /**< Used in numpy.ones() impl */
+    DPNP_FN_ONES_LIKE,    /**< Used in numpy.ones_like() impl */
+    DPNP_FN_RNG_BETA,     /**< Used in numpy.random.beta() impl  */
+    DPNP_FN_RNG_BETA_EXT, /**< Used in numpy.random.beta() impl, requires extra
+                             parameters */
+    DPNP_FN_RNG_BINOMIAL, /**< Used in numpy.random.binomial() impl  */
     DPNP_FN_RNG_BINOMIAL_EXT,  /**< Used in numpy.random.binomial() impl,
                                   requires extra parameters */
     DPNP_FN_RNG_CHISQUARE,     /**< Used in numpy.random.chisquare() impl  */

dpnp/backend/src/dpnp_fptr.hpp

@@ -161,6 +161,5 @@ class dpnp_less_comp
 void func_map_init_arraycreation(func_map_t &fmap);
 void func_map_init_linalg(func_map_t &fmap);
 void func_map_init_random(func_map_t &fmap);
-void func_map_init_sorting(func_map_t &fmap);
 
 #endif // BACKEND_FPTR_H

dpnp/backend/src/dpnp_iface_fptr.cpp

@@ -101,7 +101,6 @@ static func_map_t func_map_init()
     func_map_init_arraycreation(fmap);
     func_map_init_linalg(fmap);
     func_map_init_random(fmap);
-    func_map_init_sorting(fmap);
 
     return fmap;
 };

dpnp/dpnp_algo/CMakeLists.txt

@@ -27,13 +27,8 @@
 # THE POSSIBILITY OF SUCH DAMAGE.
 # *****************************************************************************
 
-set(dpnp_algo_pyx_deps ${CMAKE_CURRENT_SOURCE_DIR}/dpnp_algo_sorting.pxi)
-
 build_dpnp_cython_ext_with_backend(
   dpnp_algo
   ${CMAKE_CURRENT_SOURCE_DIR}/dpnp_algo.pyx
   dpnp/dpnp_algo
 )
-
-add_custom_target(_dpnp_algo_deps DEPENDS ${dpnp_algo_pyx_deps})
-add_dependencies(dpnp_algo _dpnp_algo_deps)

dpnp/dpnp_algo/dpnp_algo.pxd

@@ -32,7 +32,6 @@
 # need this namespace for Enum import
 cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName":
     cdef enum DPNPFuncName "DPNPFuncName":
-        DPNP_FN_PARTITION_EXT
         DPNP_FN_RNG_BETA_EXT
         DPNP_FN_RNG_BINOMIAL_EXT
         DPNP_FN_RNG_CHISQUARE_EXT

dpnp/dpnp_algo/dpnp_algo.pyx

@@ -40,18 +40,12 @@ and the rest of the library
 
 import dpnp
 
-cimport dpctl as c_dpctl
-
 cimport dpnp.dpnp_utils as utils
-from dpnp.dpnp_algo cimport shape_elem_type, shape_type_c
 
 __all__ = [
 ]
 
 
-include "dpnp_algo_sorting.pxi"
-
-
 """
 Internal functions
 """