Move ti.argmax()/argmin() to dpctl_ext.tensor and reuse them in dpnp

Author: vlad-perevezentsev

dpctl_ext/tensor/__init__.py

@@ -84,6 +84,8 @@
 from ._accumulation import cumulative_logsumexp, cumulative_prod, cumulative_sum
 from ._clip import clip
 from ._reduction import (
+    argmax,
+    argmin,
     max,
     min,
 )
@@ -102,6 +104,8 @@
     "all",
     "any",
     "arange",
+    "argmax",
+    "argmin",
     "argsort",
     "asarray",
     "asnumpy",

dpctl_ext/tensor/_reduction.py

@@ -137,72 +137,6 @@ def _comparison_over_axis(x, axis, keepdims, out, _reduction_fn):
     return out
 
 
-def max(x, /, *, axis=None, keepdims=False, out=None):
-    """
-    Calculates the maximum value of the input array ``x``.
-
-    Args:
-        x (usm_ndarray):
-            input array.
-        axis (Optional[int, Tuple[int, ...]]):
-            axis or axes along which maxima must be computed. If a tuple
-            of unique integers, the maxima are computed over multiple axes.
-            If ``None``, the max is computed over the entire array.
-            Default: ``None``.
-        keepdims (Optional[bool]):
-            if ``True``, the reduced axes (dimensions) are included in the
-            result as singleton dimensions, so that the returned array remains
-            compatible with the input arrays according to Array Broadcasting
-            rules. Otherwise, if ``False``, the reduced axes are not included
-            in the returned array. Default: ``False``.
-        out (Optional[usm_ndarray]):
-            the array into which the result is written.
-            The data type of ``out`` must match the expected shape and the
-            expected data type of the result.
-            If ``None`` then a new array is returned. Default: ``None``.
-
-    Returns:
-        usm_ndarray:
-            an array containing the maxima. If the max was computed over the
-            entire array, a zero-dimensional array is returned. The returned
-            array has the same data type as ``x``.
-    """
-    return _comparison_over_axis(x, axis, keepdims, out, tri._max_over_axis)
-
-
-def min(x, /, *, axis=None, keepdims=False, out=None):
-    """
-    Calculates the minimum value of the input array ``x``.
-
-    Args:
-        x (usm_ndarray):
-            input array.
-        axis (Optional[int, Tuple[int, ...]]):
-            axis or axes along which minima must be computed. If a tuple
-            of unique integers, the minima are computed over multiple axes.
-            If ``None``, the min is computed over the entire array.
-            Default: ``None``.
-        keepdims (Optional[bool]):
-            if ``True``, the reduced axes (dimensions) are included in the
-            result as singleton dimensions, so that the returned array remains
-            compatible with the input arrays according to Array Broadcasting
-            rules. Otherwise, if ``False``, the reduced axes are not included
-            in the returned array. Default: ``False``.
-        out (Optional[usm_ndarray]):
-            the array into which the result is written.
-            The data type of ``out`` must match the expected shape and the
-            expected data type of the result.
-            If ``None`` then a new array is returned. Default: ``None``.
-
-    Returns:
-        usm_ndarray:
-            an array containing the minima. If the min was computed over the
-            entire array, a zero-dimensional array is returned. The returned
-            array has the same data type as ``x``.
-    """
-    return _comparison_over_axis(x, axis, keepdims, out, tri._min_over_axis)
-
-
 def _search_over_axis(x, axis, keepdims, out, _reduction_fn):
     if not isinstance(x, dpt.usm_ndarray):
         raise TypeError(f"Expected dpctl.tensor.usm_ndarray, got {type(x)}")
@@ -376,18 +310,17 @@ def argmin(x, /, *, axis=None, keepdims=False, out=None):
     return _search_over_axis(x, axis, keepdims, out, tri._argmin_over_axis)
 
 
-def count_nonzero(x, /, *, axis=None, keepdims=False, out=None):
+def max(x, /, *, axis=None, keepdims=False, out=None):
     """
-    Counts the number of elements in the input array ``x`` which are non-zero.
+    Calculates the maximum value of the input array ``x``.
 
     Args:
         x (usm_ndarray):
             input array.
         axis (Optional[int, Tuple[int, ...]]):
-            axis or axes along which to count. If a tuple of unique integers,
-            the number of non-zero values are computed over multiple axes.
-            If ``None``, the number of non-zero values is computed over the
-            entire array.
+            axis or axes along which maxima must be computed. If a tuple
+            of unique integers, the maxima are computed over multiple axes.
+            If ``None``, the max is computed over the entire array.
             Default: ``None``.
         keepdims (Optional[bool]):
             if ``True``, the reduced axes (dimensions) are included in the
@@ -397,23 +330,47 @@ def count_nonzero(x, /, *, axis=None, keepdims=False, out=None):
             in the returned array. Default: ``False``.
         out (Optional[usm_ndarray]):
             the array into which the result is written.
-            The data type of ``out`` must match the expected shape and data
-            type.
+            The data type of ``out`` must match the expected shape and the
+            expected data type of the result.
             If ``None`` then a new array is returned. Default: ``None``.
 
     Returns:
         usm_ndarray:
-            an array containing the count of non-zero values. If the sum was
-            computed over the entire array, a zero-dimensional array is
-            returned. The returned array will have the default array index data
-            type.
+            an array containing the maxima. If the max was computed over the
+            entire array, a zero-dimensional array is returned. The returned
+            array has the same data type as ``x``.
     """
-    if x.dtype != dpt.bool:
-        x = dpt_ext.astype(x, dpt.bool, copy=False)
-    return sum(
-        x,
-        axis=axis,
-        dtype=ti.default_device_index_type(x.sycl_device),
-        keepdims=keepdims,
-        out=out,
-    )
+    return _comparison_over_axis(x, axis, keepdims, out, tri._max_over_axis)
+
+
+def min(x, /, *, axis=None, keepdims=False, out=None):
+    """
+    Calculates the minimum value of the input array ``x``.
+
+    Args:
+        x (usm_ndarray):
+            input array.
+        axis (Optional[int, Tuple[int, ...]]):
+            axis or axes along which minima must be computed. If a tuple
+            of unique integers, the minima are computed over multiple axes.
+            If ``None``, the min is computed over the entire array.
+            Default: ``None``.
+        keepdims (Optional[bool]):
+            if ``True``, the reduced axes (dimensions) are included in the
+            result as singleton dimensions, so that the returned array remains
+            compatible with the input arrays according to Array Broadcasting
+            rules. Otherwise, if ``False``, the reduced axes are not included
+            in the returned array. Default: ``False``.
+        out (Optional[usm_ndarray]):
+            the array into which the result is written.
+            The data type of ``out`` must match the expected shape and the
+            expected data type of the result.
+            If ``None`` then a new array is returned. Default: ``None``.
+
+    Returns:
+        usm_ndarray:
+            an array containing the minima. If the min was computed over the
+            entire array, a zero-dimensional array is returned. The returned
+            array has the same data type as ``x``.
+    """
+    return _comparison_over_axis(x, axis, keepdims, out, tri._min_over_axis)

dpnp/dpnp_iface_searching.py

@@ -39,12 +39,10 @@
 
 """
 
-import dpctl.tensor as dpt
-
 # pylint: disable=no-name-in-module
 # TODO: revert to `import dpctl.tensor...`
 # when dpnp fully migrates dpctl/tensor
-import dpctl_ext.tensor as dpt_ext
+import dpctl_ext.tensor as dpt
 import dpctl_ext.tensor._tensor_impl as dti
 import dpnp
 
@@ -376,13 +374,13 @@ def searchsorted(a, v, side="left", sorter=None):
 
     usm_a = dpnp.get_usm_ndarray(a)
     if dpnp.isscalar(v):
-        usm_v = dpt_ext.asarray(v, sycl_queue=a.sycl_queue, usm_type=a.usm_type)
+        usm_v = dpt.asarray(v, sycl_queue=a.sycl_queue, usm_type=a.usm_type)
     else:
         usm_v = dpnp.get_usm_ndarray(v)
 
     usm_sorter = None if sorter is None else dpnp.get_usm_ndarray(sorter)
     return dpnp_array._create_from_usm_ndarray(
-        dpt_ext.searchsorted(usm_a, usm_v, side=side, sorter=usm_sorter)
+        dpt.searchsorted(usm_a, usm_v, side=side, sorter=usm_sorter)
     )
 
 
@@ -474,7 +472,5 @@ def where(condition, x=None, y=None, /, *, order="K", out=None):
     usm_condition = dpnp.get_usm_ndarray(condition)
 
     usm_out = None if out is None else dpnp.get_usm_ndarray(out)
-    usm_res = dpt_ext.where(
-        usm_condition, usm_x, usm_y, order=order, out=usm_out
-    )
+    usm_res = dpt.where(usm_condition, usm_x, usm_y, order=order, out=usm_out)
     return dpnp.get_result_array(usm_res, out)