Separate wrapping function to reduction utils

Author: antonwolfy

dpnp/dpnp_iface_mathematical.py

@@ -55,10 +55,8 @@
 
 import dpnp
 import dpnp.backend.extensions.vm._vm_impl as vmi
-from dpnp.backend.extensions.sycl_ext import _sycl_ext_impl
-from dpnp.dpnp_array import dpnp_array
-from dpnp.dpnp_utils import call_origin, get_usm_allocations
 
+from .backend.extensions.sycl_ext import _sycl_ext_impl
 from .dpnp_algo import (
     dpnp_cumprod,
     dpnp_ediff1d,
@@ -81,7 +79,10 @@
     acceptance_fn_sign,
     acceptance_fn_subtract,
 )
+from .dpnp_array import dpnp_array
+from .dpnp_utils import call_origin, get_usm_allocations
 from .dpnp_utils.dpnp_utils_linearalgebra import dpnp_cross
+from .dpnp_utils.dpnp_utils_reduction import dpnp_wrap_reduction_call
 
 __all__ = [
     "abs",
@@ -158,33 +159,14 @@ def _append_to_diff_array(a, axis, combined, values):
     combined.append(values)
 
 
-def _wrap_reduction_call(a, dtype, out, _reduction_fn, *args, **kwargs):
-    """Wrap a call of reduction functions from dpctl.tensor interface."""
+def _get_reduction_res_dt(a, dtype, _out):
+    """Get a data type used by dpctl for result array in reduction function."""
 
-    input_out = out
-    if out is None:
-        usm_out = None
-    else:
-        dpnp.check_supported_arrays_type(out)
-
-        # get a data type used by dpctl for result array in reduction function
-        if dtype is None:
-            res_dt = dtu._default_accumulation_dtype(a.dtype, a.sycl_queue)
-        else:
-            res_dt = dpnp.dtype(dtype)
-            res_dt = dtu._to_device_supported_dtype(res_dt, a.sycl_device)
+    if dtype is None:
+        return dtu._default_accumulation_dtype(a.dtype, a.sycl_queue)
 
-        # dpctl requires strict data type matching of out array with the result
-        if out.dtype != res_dt:
-            out = dpnp.astype(out, dtype=res_dt, copy=False)
-
-        usm_out = dpnp.get_usm_ndarray(out)
-
-    kwargs["dtype"] = dtype
-    kwargs["out"] = usm_out
-    res_usm = _reduction_fn(*args, **kwargs)
-    res = dpnp_array._create_from_usm_ndarray(res_usm)
-    return dpnp.get_result_array(res, input_out, casting="unsafe")
+    dtype = dpnp.dtype(dtype)
+    return dtu._to_device_supported_dtype(dtype, a.sycl_device)
 
 
 _ABS_DOCSTRING = """
@@ -868,19 +850,22 @@ def cumsum(a, axis=None, dtype=None, out=None):
     ----------
     a : {dpnp.ndarray, usm_ndarray}
         Input array.
-    axis : int, optional
-        Axis along which the cumulative sum is computed. The default (``None``)
-        is to compute the cumulative sum over the flattened array.
+    axis : {int}, optional
+        Axis along which the cumulative sum is computed. It defaults to compute
+        the cumulative sum over the flattened array.
+        Default: ``None``.
     dtype : {None, dtype}, optional
         Type of the returned array and of the accumulator in which the elements
         are summed. If `dtype` is not specified, it defaults to the dtype of
         `a`, unless `a` has an integer dtype with a precision less than that of
         the default platform integer. In that case, the default platform
         integer is used.
-    out : {dpnp.ndarray, usm_ndarray}, optional
+        Default: ``None``.
+    out : {None, dpnp.ndarray, usm_ndarray}, optional
         Alternative output array in which to place the result. It must have the
         same shape and buffer length as the expected output but the type will
         be cast if necessary.
+        Default: ``None``.
 
     Returns
     -------
@@ -930,8 +915,14 @@ def cumsum(a, axis=None, dtype=None, out=None):
     else:
         usm_a = dpnp.get_usm_ndarray(a)
 
-    return _wrap_reduction_call(
-        a, dtype, out, dpt.cumulative_sum, usm_a, axis=axis
+    return dpnp_wrap_reduction_call(
+        a,
+        out,
+        dpt.cumulative_sum,
+        _get_reduction_res_dt,
+        usm_a,
+        axis=axis,
+        dtype=dtype,
     )
 
 
@@ -945,13 +936,13 @@ def diff(a, n=1, axis=-1, prepend=None, append=None):
     ----------
     a : {dpnp.ndarray, usm_ndarray}
         Input array
-    n : int, optional
+    n : {int}, optional
         The number of times the values differ. If ``zero``, the input
         is returned as-is.
-    axis : int, optional
+    axis : {int}, optional
         The axis along which the difference is taken, default is the
         last axis.
-    prepend, append : {scalar, dpnp.ndarray, usm_ndarray}, optional
+    prepend, append : {None, scalar, dpnp.ndarray, usm_ndarray}, optional
         Values to prepend or append to `a` along axis prior to
         performing the difference. Scalar values are expanded to
         arrays with length 1 in the direction of axis and the shape
@@ -2332,8 +2323,15 @@ def prod(
     dpnp.check_limitations(initial=initial, where=where)
     usm_a = dpnp.get_usm_ndarray(a)
 
-    return _wrap_reduction_call(
-        a, dtype, out, dpt.prod, usm_a, axis=axis, keepdims=keepdims
+    return dpnp_wrap_reduction_call(
+        a,
+        out,
+        dpt.prod,
+        _get_reduction_res_dt,
+        usm_a,
+        axis=axis,
+        dtype=dtype,
+        keepdims=keepdims,
     )
 
 
@@ -2912,8 +2910,15 @@ def sum(
             return result
 
     usm_a = dpnp.get_usm_ndarray(a)
-    return _wrap_reduction_call(
-        a, dtype, out, dpt.sum, usm_a, axis=axis, keepdims=keepdims
+    return dpnp_wrap_reduction_call(
+        a,
+        out,
+        dpt.sum,
+        _get_reduction_res_dt,
+        usm_a,
+        axis=axis,
+        dtype=dtype,
+        keepdims=keepdims,
     )
 
 

dpnp/dpnp_iface_nanfunctions.py

@@ -301,7 +301,7 @@ def nancumsum(a, axis=None, dtype=None, out=None):
     ----------
     a : {dpnp.ndarray, usm_ndarray}
         Input array.
-    axis : int, optional
+    axis : {int}, optional
         Axis along which the cumulative sum is computed. The default (``None``)
         is to compute the cumulative sum over the flattened array.
     dtype : {None, dtype}, optional
@@ -310,7 +310,7 @@ def nancumsum(a, axis=None, dtype=None, out=None):
         `a`, unless `a` has an integer dtype with a precision less than that of
         the default platform integer. In that case, the default platform
         integer is used.
-    out : {dpnp.ndarray, usm_ndarray}, optional
+    out : {None, dpnp.ndarray, usm_ndarray}, optional
         Alternative output array in which to place the result. It must have the
         same shape and buffer length as the expected output but the type will
         be cast if necessary.

dpnp/dpnp_iface_searching.py

@@ -49,39 +49,24 @@
 from .dpnp_utils import (
     get_usm_allocations,
 )
+from .dpnp_utils.dpnp_utils_reduction import dpnp_wrap_reduction_call
 
 __all__ = ["argmax", "argmin", "searchsorted", "where"]
 
 
-def _wrap_search_call(a, out, _search_fn, *args, **kwargs):
-    """Wrap a call of search functions from dpctl.tensor interface."""
+def _get_search_res_dt(a, _dtype, out):
+    """Get a data type used by dpctl for result array in search function."""
 
-    input_out = out
-    if out is None:
-        usm_out = None
-    else:
-        dpnp.check_supported_arrays_type(out)
-
-        # get a data type used by dpctl for result array in search function
-        res_dt = dti.default_device_index_type(a.sycl_device)
-
-        # numpy raises TypeError if "out" data type mismatch default index type
-        if not dpnp.can_cast(out.dtype, res_dt, casting="safe"):
-            raise TypeError(
-                f"Cannot cast from {out.dtype} to {res_dt} "
-                "according to the rule safe."
-            )
-
-        # dpctl requires strict data type matching of out array with the result
-        if out.dtype != res_dt:
-            out = dpnp.astype(out, dtype=res_dt, copy=False)
-
-        usm_out = dpnp.get_usm_ndarray(out)
+    # get a data type used by dpctl for result array in search function
+    res_dt = dti.default_device_index_type(a.sycl_device)
 
-    kwargs["out"] = usm_out
-    res_usm = _search_fn(*args, **kwargs)
-    res = dpnp_array._create_from_usm_ndarray(res_usm)
-    return dpnp.get_result_array(res, input_out, casting="unsafe")
+    # numpy raises TypeError if "out" data type mismatch default index type
+    if not dpnp.can_cast(out.dtype, res_dt, casting="safe"):
+        raise TypeError(
+            f"Cannot cast from {out.dtype} to {res_dt} "
+            "according to the rule safe."
+        )
+    return res_dt
 
 
 def argmax(a, axis=None, out=None, *, keepdims=False):
@@ -163,8 +148,14 @@ def argmax(a, axis=None, out=None, *, keepdims=False):
     """
 
     usm_a = dpnp.get_usm_ndarray(a)
-    return _wrap_search_call(
-        a, out, dpt.argmax, usm_a, axis=axis, keepdims=keepdims
+    return dpnp_wrap_reduction_call(
+        a,
+        out,
+        dpt.argmax,
+        _get_search_res_dt,
+        usm_a,
+        axis=axis,
+        keepdims=keepdims,
     )
 
 
@@ -248,8 +239,14 @@ def argmin(a, axis=None, out=None, *, keepdims=False):
     """
 
     usm_a = dpnp.get_usm_ndarray(a)
-    return _wrap_search_call(
-        a, out, dpt.argmin, usm_a, axis=axis, keepdims=keepdims
+    return dpnp_wrap_reduction_call(
+        a,
+        out,
+        dpt.argmin,
+        _get_search_res_dt,
+        usm_a,
+        axis=axis,
+        keepdims=keepdims,
     )
 
 

dpnp/dpnp_iface_statistics.py

@@ -54,6 +54,7 @@
     call_origin,
     get_usm_allocations,
 )
+from .dpnp_utils.dpnp_utils_reduction import dpnp_wrap_reduction_call
 from .dpnp_utils.dpnp_utils_statistics import (
     dpnp_cov,
 )
@@ -118,28 +119,10 @@ def _count_reduce_items(arr, axis, where=True):
     return items
 
 
-def _wrap_comparison_call(a, out, _comparison_fn, *args, **kwargs):
-    """Wrap a call of comparison functions from dpctl.tensor interface."""
+def _get_comparison_res_dt(a, _dtype, _out):
+    """Get a data type used by dpctl for result array in comparison function."""
 
-    input_out = out
-    if out is None:
-        usm_out = None
-    else:
-        dpnp.check_supported_arrays_type(out)
-
-        # get dtype used by dpctl for result array in comparison function
-        res_dt = a.dtype
-
-        # dpctl requires strict data type matching of out array with the result
-        if out.dtype != res_dt:
-            out = dpnp.astype(out, dtype=res_dt, copy=False)
-
-        usm_out = dpnp.get_usm_ndarray(out)
-
-    kwargs["out"] = usm_out
-    res_usm = _comparison_fn(*args, **kwargs)
-    res = dpnp_array._create_from_usm_ndarray(res_usm)
-    return dpnp.get_result_array(res, input_out, casting="unsafe")
+    return a.dtype
 
 
 def amax(a, axis=None, out=None, keepdims=False, initial=None, where=True):
@@ -548,8 +531,14 @@ def max(a, axis=None, out=None, keepdims=False, initial=None, where=True):
     dpnp.check_limitations(initial=initial, where=where)
     usm_a = dpnp.get_usm_ndarray(a)
 
-    return _wrap_comparison_call(
-        a, out, dpt.max, usm_a, axis=axis, keepdims=keepdims
+    return dpnp_wrap_reduction_call(
+        a,
+        out,
+        dpt.max,
+        _get_comparison_res_dt,
+        usm_a,
+        axis=axis,
+        keepdims=keepdims,
     )
 
 
@@ -757,8 +746,14 @@ def min(a, axis=None, out=None, keepdims=False, initial=None, where=True):
     dpnp.check_limitations(initial=initial, where=where)
     usm_a = dpnp.get_usm_ndarray(a)
 
-    return _wrap_comparison_call(
-        a, out, dpt.min, usm_a, axis=axis, keepdims=keepdims
+    return dpnp_wrap_reduction_call(
+        a,
+        out,
+        dpt.min,
+        _get_comparison_res_dt,
+        usm_a,
+        axis=axis,
+        keepdims=keepdims,
     )