Implement kernels for median and paritition functions

CHANGELOG.md

@@ -12,6 +12,7 @@ This release drops support for Python 3.10, making Python 3.11 the minimum requi
 
 * 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
 
@@ -23,6 +24,7 @@ This release drops support for Python 3.10, making Python 3.11 the minimum requi
 * Improved performance of `dpnp.scipy.sparse.csr_matrix.sort_indices` by sorting once on a combined row and column key instead of in two passes [#3134](https://github.com/IntelPython/dpnp/pull/3134)
 * Reduced the overhead of `dpnp.ndarray.flags` and `dpnp.tensor.usm_ndarray.flags` [#3168](https://github.com/IntelPython/dpnp/pull/3168)
 * Changed `dpnp.linalg.LinAlgError` to derive from `numpy.linalg.LinAlgError`, so `except numpy.linalg.LinAlgError` also catches errors raised by dpnp linear algebra functions [#3184](https://github.com/IntelPython/dpnp/pull/3184)
+* Changed `dpnp.median`, `dpnp.nanmedian` and `dpnp.partition` with a single `kth` 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
 
@@ -36,6 +38,7 @@ This release drops support for Python 3.10, making Python 3.11 the minimum requi
 * 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)
 * Dropped support for Python 3.10 [#3173](https://github.com/IntelPython/dpnp/pull/3173)
 * Removed the `USE_GROUP_LOAD_STORE` macro from the libtensor header `utils/sycl_utils.hpp`, so `sub_group_load` and `sub_group_store` always use the `sycl_ext_oneapi_group_load_store` extension; the `sycl::sub_group` member load/store fallback it selected does not compile with DPC++ 2026.0, the minimum supported compiler [#3174](https://github.com/IntelPython/dpnp/pull/3174)
+* Removed the legacy implementation of `dpnp.partition` from the backend [#3138](https://github.com/IntelPython/dpnp/pull/3138)
 
 ### Fixed
 
@@ -44,6 +47,9 @@ This release drops support for Python 3.10, making Python 3.11 the minimum requi
 * 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 the example of handling duplicate indices in the `dpnp.tensor.put` docstring, which imported `dpctl.tensor`, did not parse and lost writes when `x` is longer than `indices` [#3165](https://github.com/IntelPython/dpnp/pull/3165)
 * Fixed a segfault when setting `writable` on a `Flags` object created without an array [#3168](https://github.com/IntelPython/dpnp/pull/3168)
+* 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)
+* Fixed `dpnp.partition` to accept NumPy integers and integer arrays as `kth`, and to raise as NumPy does for a boolean or an empty `kth` [#3138](https://github.com/IntelPython/dpnp/pull/3138)
+* Fixed `dpnp.nanmedian` raising `TypeError` for an `out` array of a lower precision data type instead of casting to it as `dpnp.median` does [#3138](https://github.com/IntelPython/dpnp/pull/3138)
 
 ### Security
 

dpnp/AGENTS.md

@@ -6,7 +6,7 @@
 - `dpnp_algo/dpnp_elementwise_common.py`: ufunc wrapper classes (`DPNPUnaryFunc`, `DPNPBinaryFunc`, two-output and specialised variants)
 - `dpnp_utils/`: shared helpers (`get_usm_allocations`, `map_dtype_to_device`, reduction wrappers)
 - `linalg/`, `fft/`, `random/`, `scipy/`, `lib/`: subpackages, each exporting through its own `__init__.py`
-- Legacy, don't extend: the Cython `dpnp_algo` path and `dpnp_backend_c` (still used by `dpnp.partition` and `dpnp.random`), and the NumPy fallback `call_origin`. New kernels go into `dpnp/libtensor` (see `dpnp/libtensor/AGENTS.md`).
+- Legacy, don't extend: the Cython `dpnp_algo` path and `dpnp_backend_c` (still used by `dpnp.random`), and the NumPy fallback `call_origin`. New kernels go into `dpnp/libtensor` (see `dpnp/libtensor/AGENTS.md`).
 
 ## Adding a Public Function
 1. Implement it in the matching `dpnp_iface_*.py` or subpackage module. Iface modules have no `__all__`.

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/AGENTS.md

@@ -3,5 +3,5 @@
 Legacy, don't extend: new kernels and bindings go into `dpnp/libtensor` (see `dpnp/libtensor/AGENTS.md`, which also holds the C++ rules and build options).
 
 ## Layout
-- `src/`, `include/` and `kernels/dpnp_krnl_*.cpp` build `dpnp_backend_c` (still used by `dpnp.partition` and `dpnp.random`); `tests/` (GTest, `-DDPNP_BACKEND_TESTS=ON`) and `examples/` are not built in CI
+- `src/`, `include/` and `kernels/dpnp_krnl_*.cpp` build `dpnp_backend_c` (still used by `dpnp.random`); `tests/` (GTest, `-DDPNP_BACKEND_TESTS=ON`) and `examples/` are not built in CI
 - `cmake/Modules/`: vendored `IntelSYCL`, `MKL`, `TBB` and `oneDPL` config modules, used by the whole build

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

@@ -72,95 +72,11 @@ const DPNPFuncType eft_C64 = DPNPFuncType::DPNP_FT_CMPLX64;
 const DPNPFuncType eft_C128 = DPNPFuncType::DPNP_FT_CMPLX128;
 const DPNPFuncType eft_BLN = DPNPFuncType::DPNP_FT_BOOL;
 
-/**
- * Implements std::is_same<> with variadic number of types to compare with
- * and when type T has to match only one of types Ts.
- */
-template <typename T, typename... Ts>
-struct is_any : std::disjunction<std::is_same<T, Ts>...>
-{
-};
-
-/**
- * Implements std::is_same<> with variadic number of types to compare with
- * and when type T has to match every type from Ts sequence.
- */
-template <typename T, typename... Ts>
-struct are_same : std::conjunction<std::is_same<T, Ts>...>
-{
-};
-
-/**
- * A template constant to check if type T matches any type from Ts.
- */
-template <typename T, typename... Ts>
-constexpr auto is_any_v = is_any<T, Ts...>::value;
-
-/**
- * A template constat to check if both types T1 and T2 match every type from Ts
- * sequence.
- */
-template <typename T1, typename T2, typename... Ts>
-constexpr auto both_types_are_same =
-    std::conjunction_v<is_any<T1, Ts...>, are_same<T1, T2>>;
-
-/**
- * @brief If the type _Tp is a reference type, provides the member typedef type
- * which is the type referred to by _Tp with its topmost cv-qualifiers removed.
- * Otherwise type is _Tp with its topmost cv-qualifiers removed.
- *
- * @note std::remove_cvref is only available since c++20
- */
-template <typename _Tp>
-using dpnp_remove_cvref_t =
-    typename std::remove_cv_t<typename std::remove_reference_t<_Tp>>;
-
-/**
- * @brief "<" comparison with complex types support.
- *
- * @note return a result of lexicographical "<" comparison for complex types.
- */
-class dpnp_less_comp
-{
-public:
-    template <typename _Xp, typename _Yp>
-    bool operator()(_Xp &&__x, _Yp &&__y) const
-    {
-        if constexpr (both_types_are_same<
-                          dpnp_remove_cvref_t<_Xp>, dpnp_remove_cvref_t<_Yp>,
-                          std::complex<float>, std::complex<double>>) {
-            bool ret = false;
-            _Xp a = std::forward<_Xp>(__x);
-            _Yp b = std::forward<_Yp>(__y);
-
-            if (a.real() < b.real()) {
-                ret = (a.imag() == a.imag() || b.imag() != b.imag());
-            }
-            else if (a.real() > b.real()) {
-                ret = (b.imag() != b.imag() && a.imag() == a.imag());
-            }
-            else if (a.real() == b.real() ||
-                     (a.real() != a.real() && b.real() != b.real())) {
-                ret = (a.imag() < b.imag() ||
-                       (b.imag() != b.imag() && a.imag() == a.imag()));
-            }
-            else {
-                ret = (b.real() != b.real());
-            }
-            return ret;
-        }
-        else {
-            return std::forward<_Xp>(__x) < std::forward<_Yp>(__y);
-        }
-    }
-};
-
 /**
  * FPTR interface initialization functions
  */
 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