Improve dpnp.partition implementation (#2766)

The PR propose to improve implementation and to use `dpnp.sort` call
when
- input array has number of dimensions > 1
- input array has previously not supported integer dtype
- `axis` keyword is passed (previously not supported)
- sequence of `kth` is passed (previously not supported)
In case of `ndim > 1` previously the implementation from legacy backend
was used, which is significantly slow (see performance comparation
below). It used a copy of input data into the shared USM memory and
included computations on the host.

This PR proposes to reuse `dpnp.sort` for all the above cases.
While in case when the legacy implementation is stable and fast (for 1D
input array), it will remain, because it relays on `std::nth_element`
from OneDPL.

The benchmark results were collected on PVC with help of the below code:
```python
import dpnp, numpy as np
from dpnp.tests.helper import generate_random_numpy_array

a = generate_random_numpy_array(10**7, dtype=np.float64, seed_value=117)
ia = dpnp.array(a)
%timeit x = dpnp.partition(ia, 513); x.sycl_queue.wait()
```

Below tables contains data in case of 1D input array (shape=(10**7,)),
where the implementation path was kept the same, plus adding support of
missing integer dtypes using fallback on the sort function:
| Implementation | int32 | uint32 | int64 | uint64 | float32 | float64 |
complex64 | complex128 |

|--------|--------|--------|--------|--------|--------|--------|--------|--------|
| old (legacy backend) | 7.46 ms | not supported | 9.46 ms | not
supported | 7.39 ms | 8.92 ms | 10.9 ms | 21.2 ms |
| new (backend + sort) | 7.34 ms | 10.8 ms | 9.48 ms | 12.5 ms | 7.37 ms
| 8.89 ms | 11 ms | 21.2 ms |

The following code was used for 2D input array with shape=(10**4,
10**4):
```python
import dpnp, numpy as np
from dpnp.tests.helper import generate_random_numpy_array

a = generate_random_numpy_array((10**4, 10**4), dtype=np.float64, seed_value=117)
ia = dpnp.array(a)
%timeit x = dpnp.partition(ia, 1513); x.sycl_queue.wait()
```

In that case the new implementation is fully based on the sort call:
| Implementation | int32 | int64 | float32 | float64 | complex64 |
complex128 |
|--------|--------|--------|--------|--------|--------|--------|
| old (legacy backend) | 6.4 s | 6.89 s | 7.36 s | 7.66 s | 8.61 s | 10
s |
| new (sort) | 57.4 ms | 64.7 ms | 62.2 ms | 68 ms | 77 ms | 151 ms | 9c4aed2

Author: github-actions[bot]

.buildinfo

@@ -1,4 +1,4 @@
 # Sphinx build info version 1
 # This file records the configuration used when building these files. When it is not found, a full rebuild will be done.
-config: dabfa318cc06c6ce826220cdde96bd82
+config: bf51692a0db2c4243b085a8d82c6bc8b
 tags: 645f666f9bcd5a90fca523b33c5a78b7

_modules/dpnp/dpnp_array.html

@@ -14,14 +14,14 @@
     </script>
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>dpnp.dpnp_array &mdash; Data Parallel Extension for NumPy 0.20.0dev2+20.g9d6d5a50e6f documentation</title>
+  <title>dpnp.dpnp_array &mdash; Data Parallel Extension for NumPy 0.20.0dev2+21.g9c4aed22eb8 documentation</title>
       <link rel="stylesheet" type="text/css" href="../../_static/pygments.css?v=03e43079" />
       <link rel="stylesheet" type="text/css" href="../../_static/css/theme.css?v=9edc463e" />
 
 
       <script src="../../_static/jquery.js?v=5d32c60e"></script>
       <script src="../../_static/_sphinx_javascript_frameworks_compat.js?v=2cd50e6c"></script>
-      <script src="../../_static/documentation_options.js?v=3ba98b21"></script>
+      <script src="../../_static/documentation_options.js?v=20ce9460"></script>
       <script src="../../_static/doctools.js?v=9bcbadda"></script>
       <script src="../../_static/sphinx_highlight.js?v=dc90522c"></script>
     <script src="../../_static/js/theme.js"></script>
@@ -1623,35 +1623,54 @@ <h1>Source code for dpnp.dpnp_array</h1><div class="highlight"><pre>
 <a class="viewcode-back" href="../../reference/generated/dpnp.dpnp_array.dpnp_array.partition.html#dpnp.dpnp_array.dpnp_array.partition">[docs]</a>
     <span class="k">def</span><span class="w"> </span><span class="nf">partition</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">/</span><span class="p">,</span> <span class="n">kth</span><span class="p">,</span> <span class="n">axis</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span> <span class="n">kind</span><span class="o">=</span><span class="s2">&quot;introselect&quot;</span><span class="p">,</span> <span class="n">order</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
 <span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
-<span class="sd">        Return a partitioned copy of an array.</span>
+<span class="sd">        Partially sorts the elements in the array in such a way that the value</span>
+<span class="sd">        of the element in k-th position is in the position it would be in a</span>
+<span class="sd">        sorted array. In the output array, all elements smaller than the k-th</span>
+<span class="sd">        element are located to the left of this element and all equal or</span>
+<span class="sd">        greater are located to its right. The ordering of the elements in the</span>
+<span class="sd">        two partitions on the either side of the k-th element in the output</span>
+<span class="sd">        array is undefined.</span>
 
-<span class="sd">        Rearranges the elements in the array in such a way that the value of</span>
-<span class="sd">        the element in `kth` position is in the position it would be in</span>
-<span class="sd">        a sorted array.</span>
+<span class="sd">        Refer to `dpnp.partition` for full documentation.</span>
 
-<span class="sd">        All elements smaller than the `kth` element are moved before this</span>
-<span class="sd">        element and all equal or greater are moved behind it. The ordering</span>
-<span class="sd">        of the elements in the two partitions is undefined.</span>
+<span class="sd">        kth : {int, sequence of ints}</span>
+<span class="sd">            Element index to partition by. The kth element value will be in its</span>
+<span class="sd">            final sorted position and all smaller elements will be moved before</span>
+<span class="sd">            it and all equal or greater elements behind it.</span>
+<span class="sd">            The order of all elements in the partitions is undefined. If</span>
+<span class="sd">            provided with a sequence of kth it will partition all elements</span>
+<span class="sd">            indexed by kth of them into their sorted position at once.</span>
+<span class="sd">        axis : int, optional</span>
+<span class="sd">            Axis along which to sort. The default is ``-1``, which means sort</span>
+<span class="sd">            along the the last axis.</span>
 
-<span class="sd">        Refer to `dpnp.partition` for full documentation.</span>
+<span class="sd">            Default: ``-1``.</span>
 
 <span class="sd">        See Also</span>
 <span class="sd">        --------</span>
 <span class="sd">        :obj:`dpnp.partition` : Return a partitioned copy of an array.</span>
+<span class="sd">        :obj:`dpnp.argpartition` : Indirect partition.</span>
+<span class="sd">        :obj:`dpnp.sort` : Full sort.</span>
 
 <span class="sd">        Examples</span>
 <span class="sd">        --------</span>
 <span class="sd">        &gt;&gt;&gt; import dpnp as np</span>
 <span class="sd">        &gt;&gt;&gt; a = np.array([3, 4, 2, 1])</span>
 <span class="sd">        &gt;&gt;&gt; a.partition(3)</span>
 <span class="sd">        &gt;&gt;&gt; a</span>
+<span class="sd">        array([1, 2, 3, 4]) # may vary</span>
+
+<span class="sd">        &gt;&gt;&gt; a.partition((1, 3))</span>
+<span class="sd">        &gt;&gt;&gt; a</span>
 <span class="sd">        array([1, 2, 3, 4])</span>
 
 <span class="sd">        &quot;&quot;&quot;</span>
 
-        <span class="bp">self</span><span class="o">.</span><span class="n">_array_obj</span> <span class="o">=</span> <span class="n">dpnp</span><span class="o">.</span><span class="n">partition</span><span class="p">(</span>
-            <span class="bp">self</span><span class="p">,</span> <span class="n">kth</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="n">axis</span><span class="p">,</span> <span class="n">kind</span><span class="o">=</span><span class="n">kind</span><span class="p">,</span> <span class="n">order</span><span class="o">=</span><span class="n">order</span>
-        <span class="p">)</span><span class="o">.</span><span class="n">get_array</span><span class="p">()</span></div>
+        <span class="k">if</span> <span class="n">axis</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
+            <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span>
+                <span class="s2">&quot;&#39;NoneType&#39; object cannot be interpreted as an integer&quot;</span>
+            <span class="p">)</span>
+        <span class="bp">self</span><span class="p">[</span><span class="o">...</span><span class="p">]</span> <span class="o">=</span> <span class="n">dpnp</span><span class="o">.</span><span class="n">partition</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">kth</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="n">axis</span><span class="p">,</span> <span class="n">kind</span><span class="o">=</span><span class="n">kind</span><span class="p">,</span> <span class="n">order</span><span class="o">=</span><span class="n">order</span><span class="p">)</span></div>
 
 
 <div class="viewcode-block" id="dpnp_array.prod">

_modules/dpnp/dpnp_array_api_info.html

@@ -14,14 +14,14 @@
     </script>
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>dpnp.dpnp_array_api_info &mdash; Data Parallel Extension for NumPy 0.20.0dev2+20.g9d6d5a50e6f documentation</title>
+  <title>dpnp.dpnp_array_api_info &mdash; Data Parallel Extension for NumPy 0.20.0dev2+21.g9c4aed22eb8 documentation</title>
       <link rel="stylesheet" type="text/css" href="../../_static/pygments.css?v=03e43079" />
       <link rel="stylesheet" type="text/css" href="../../_static/css/theme.css?v=9edc463e" />
 
 
       <script src="../../_static/jquery.js?v=5d32c60e"></script>
       <script src="../../_static/_sphinx_javascript_frameworks_compat.js?v=2cd50e6c"></script>
-      <script src="../../_static/documentation_options.js?v=3ba98b21"></script>
+      <script src="../../_static/documentation_options.js?v=20ce9460"></script>
       <script src="../../_static/doctools.js?v=9bcbadda"></script>
       <script src="../../_static/sphinx_highlight.js?v=dc90522c"></script>
     <script src="../../_static/js/theme.js"></script>

_modules/dpnp/dpnp_flatiter.html

@@ -14,14 +14,14 @@
     </script>
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>dpnp.dpnp_flatiter &mdash; Data Parallel Extension for NumPy 0.20.0dev2+20.g9d6d5a50e6f documentation</title>
+  <title>dpnp.dpnp_flatiter &mdash; Data Parallel Extension for NumPy 0.20.0dev2+21.g9c4aed22eb8 documentation</title>
       <link rel="stylesheet" type="text/css" href="../../_static/pygments.css?v=03e43079" />
       <link rel="stylesheet" type="text/css" href="../../_static/css/theme.css?v=9edc463e" />
 
 
       <script src="../../_static/jquery.js?v=5d32c60e"></script>
       <script src="../../_static/_sphinx_javascript_frameworks_compat.js?v=2cd50e6c"></script>
-      <script src="../../_static/documentation_options.js?v=3ba98b21"></script>
+      <script src="../../_static/documentation_options.js?v=20ce9460"></script>
       <script src="../../_static/doctools.js?v=9bcbadda"></script>
       <script src="../../_static/sphinx_highlight.js?v=dc90522c"></script>
     <script src="../../_static/js/theme.js"></script>

_modules/dpnp/dpnp_iface.html

@@ -14,14 +14,14 @@
     </script>
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>dpnp.dpnp_iface &mdash; Data Parallel Extension for NumPy 0.20.0dev2+20.g9d6d5a50e6f documentation</title>
+  <title>dpnp.dpnp_iface &mdash; Data Parallel Extension for NumPy 0.20.0dev2+21.g9c4aed22eb8 documentation</title>
       <link rel="stylesheet" type="text/css" href="../../_static/pygments.css?v=03e43079" />
       <link rel="stylesheet" type="text/css" href="../../_static/css/theme.css?v=9edc463e" />
 
 
       <script src="../../_static/jquery.js?v=5d32c60e"></script>
       <script src="../../_static/_sphinx_javascript_frameworks_compat.js?v=2cd50e6c"></script>
-      <script src="../../_static/documentation_options.js?v=3ba98b21"></script>
+      <script src="../../_static/documentation_options.js?v=20ce9460"></script>
       <script src="../../_static/doctools.js?v=9bcbadda"></script>
       <script src="../../_static/sphinx_highlight.js?v=dc90522c"></script>
     <script src="../../_static/js/theme.js"></script>

_modules/dpnp/dpnp_iface_arraycreation.html

@@ -14,14 +14,14 @@
     </script>
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>dpnp.dpnp_iface_arraycreation &mdash; Data Parallel Extension for NumPy 0.20.0dev2+20.g9d6d5a50e6f documentation</title>
+  <title>dpnp.dpnp_iface_arraycreation &mdash; Data Parallel Extension for NumPy 0.20.0dev2+21.g9c4aed22eb8 documentation</title>
       <link rel="stylesheet" type="text/css" href="../../_static/pygments.css?v=03e43079" />
       <link rel="stylesheet" type="text/css" href="../../_static/css/theme.css?v=9edc463e" />
 
 
       <script src="../../_static/jquery.js?v=5d32c60e"></script>
       <script src="../../_static/_sphinx_javascript_frameworks_compat.js?v=2cd50e6c"></script>
-      <script src="../../_static/documentation_options.js?v=3ba98b21"></script>
+      <script src="../../_static/documentation_options.js?v=20ce9460"></script>
       <script src="../../_static/doctools.js?v=9bcbadda"></script>
       <script src="../../_static/sphinx_highlight.js?v=dc90522c"></script>
     <script src="../../_static/js/theme.js"></script>

_modules/dpnp/dpnp_iface_bitwise.html

@@ -14,14 +14,14 @@
     </script>
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>dpnp.dpnp_iface_bitwise &mdash; Data Parallel Extension for NumPy 0.20.0dev2+20.g9d6d5a50e6f documentation</title>
+  <title>dpnp.dpnp_iface_bitwise &mdash; Data Parallel Extension for NumPy 0.20.0dev2+21.g9c4aed22eb8 documentation</title>
       <link rel="stylesheet" type="text/css" href="../../_static/pygments.css?v=03e43079" />
       <link rel="stylesheet" type="text/css" href="../../_static/css/theme.css?v=9edc463e" />
 
 
       <script src="../../_static/jquery.js?v=5d32c60e"></script>
       <script src="../../_static/_sphinx_javascript_frameworks_compat.js?v=2cd50e6c"></script>
-      <script src="../../_static/documentation_options.js?v=3ba98b21"></script>
+      <script src="../../_static/documentation_options.js?v=20ce9460"></script>
       <script src="../../_static/doctools.js?v=9bcbadda"></script>
       <script src="../../_static/sphinx_highlight.js?v=dc90522c"></script>
     <script src="../../_static/js/theme.js"></script>

_modules/dpnp/dpnp_iface_counting.html

@@ -14,14 +14,14 @@
     </script>
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>dpnp.dpnp_iface_counting &mdash; Data Parallel Extension for NumPy 0.20.0dev2+20.g9d6d5a50e6f documentation</title>
+  <title>dpnp.dpnp_iface_counting &mdash; Data Parallel Extension for NumPy 0.20.0dev2+21.g9c4aed22eb8 documentation</title>
       <link rel="stylesheet" type="text/css" href="../../_static/pygments.css?v=03e43079" />
       <link rel="stylesheet" type="text/css" href="../../_static/css/theme.css?v=9edc463e" />
 
 
       <script src="../../_static/jquery.js?v=5d32c60e"></script>
       <script src="../../_static/_sphinx_javascript_frameworks_compat.js?v=2cd50e6c"></script>
-      <script src="../../_static/documentation_options.js?v=3ba98b21"></script>
+      <script src="../../_static/documentation_options.js?v=20ce9460"></script>
       <script src="../../_static/doctools.js?v=9bcbadda"></script>
       <script src="../../_static/sphinx_highlight.js?v=dc90522c"></script>
     <script src="../../_static/js/theme.js"></script>

_modules/dpnp/dpnp_iface_functional.html

@@ -14,14 +14,14 @@
     </script>
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>dpnp.dpnp_iface_functional &mdash; Data Parallel Extension for NumPy 0.20.0dev2+20.g9d6d5a50e6f documentation</title>
+  <title>dpnp.dpnp_iface_functional &mdash; Data Parallel Extension for NumPy 0.20.0dev2+21.g9c4aed22eb8 documentation</title>
       <link rel="stylesheet" type="text/css" href="../../_static/pygments.css?v=03e43079" />
       <link rel="stylesheet" type="text/css" href="../../_static/css/theme.css?v=9edc463e" />
 
 
       <script src="../../_static/jquery.js?v=5d32c60e"></script>
       <script src="../../_static/_sphinx_javascript_frameworks_compat.js?v=2cd50e6c"></script>
-      <script src="../../_static/documentation_options.js?v=3ba98b21"></script>
+      <script src="../../_static/documentation_options.js?v=20ce9460"></script>
       <script src="../../_static/doctools.js?v=9bcbadda"></script>
       <script src="../../_static/sphinx_highlight.js?v=dc90522c"></script>
     <script src="../../_static/js/theme.js"></script>

_modules/dpnp/dpnp_iface_histograms.html

@@ -14,14 +14,14 @@
     </script>
 
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
-  <title>dpnp.dpnp_iface_histograms &mdash; Data Parallel Extension for NumPy 0.20.0dev2+20.g9d6d5a50e6f documentation</title>
+  <title>dpnp.dpnp_iface_histograms &mdash; Data Parallel Extension for NumPy 0.20.0dev2+21.g9c4aed22eb8 documentation</title>
       <link rel="stylesheet" type="text/css" href="../../_static/pygments.css?v=03e43079" />
       <link rel="stylesheet" type="text/css" href="../../_static/css/theme.css?v=9edc463e" />
 
 
       <script src="../../_static/jquery.js?v=5d32c60e"></script>
       <script src="../../_static/_sphinx_javascript_frameworks_compat.js?v=2cd50e6c"></script>
-      <script src="../../_static/documentation_options.js?v=3ba98b21"></script>
+      <script src="../../_static/documentation_options.js?v=20ce9460"></script>
       <script src="../../_static/doctools.js?v=9bcbadda"></script>
       <script src="../../_static/sphinx_highlight.js?v=dc90522c"></script>
     <script src="../../_static/js/theme.js"></script>