Add dpnp.minmax with a fused min/max reduction kernel

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 * Added support for boolean input in `dpnp.positive` and the `+` operator of `dpnp.ndarray` [#3128](https://github.com/IntelPython/dpnp/pull/3128)
+* Added implementation of `dpnp.minmax`, computing the minimum and the maximum of an array in a single pass by a new fused reduction kernel [#3135](https://github.com/IntelPython/dpnp/pull/3135)
 
 ### Changed
 

doc/reference/math.rst

@@ -208,6 +208,8 @@ Extrema finding
    fmin
    nanmin
 
+   minmax
+
 
 Miscellaneous
 -------------

dpnp/__init__.py

@@ -547,6 +547,7 @@
     mean,
     median,
     min,
+    minmax,
     ptp,
     std,
     var,
@@ -1016,6 +1017,7 @@
     "mean",
     "median",
     "min",
+    "minmax",
     "nanmax",
     "nanmean",
     "nanmedian",

dpnp/dpnp_iface_statistics.py

@@ -50,6 +50,7 @@
 import dpnp.backend.extensions.statistics._statistics_impl as statistics_ext
 import dpnp.tensor as dpt
 import dpnp.tensor._tensor_elementwise_impl as ti
+import dpnp.tensor._tensor_impl as dti
 
 from .dpnp_utils import get_usm_allocations
 from .dpnp_utils.dpnp_utils_common import (
@@ -59,6 +60,7 @@
 from .dpnp_utils.dpnp_utils_reduction import dpnp_wrap_reduction_call
 from .dpnp_utils.dpnp_utils_statistics import dpnp_cov, dpnp_median
 from .tensor._numpy_helper import normalize_axis_index
+from .tensor._reduction import _minmax as dpt_minmax
 
 
 def _count_reduce_items(arr, axis, where=True):
@@ -1094,6 +1096,7 @@ def max(a, axis=None, out=None, keepdims=False, initial=None, where=True):
     See Also
     --------
     :obj:`dpnp.min` : Return the minimum of an array.
+    :obj:`dpnp.minmax` : Return the minimum and maximum of an array.
     :obj:`dpnp.maximum` : Element-wise maximum of two arrays, propagates NaNs.
     :obj:`dpnp.fmax` : Element-wise maximum of two arrays, ignores NaNs.
     :obj:`dpnp.nanmax` : The maximum value of an array along a given axis,
@@ -1369,6 +1372,7 @@ def min(a, axis=None, out=None, keepdims=False, initial=None, where=True):
     See Also
     --------
     :obj:`dpnp.max` : Return the maximum of an array.
+    :obj:`dpnp.minmax` : Return the minimum and maximum of an array.
     :obj:`dpnp.minimum` : Element-wise minimum of two arrays, propagates NaNs.
     :obj:`dpnp.fmin` : Element-wise minimum of two arrays, ignores NaNs.
     :obj:`dpnp.nanmin` : The minimum value of an array along a given axis,
@@ -1409,6 +1413,135 @@ def min(a, axis=None, out=None, keepdims=False, initial=None, where=True):
     )
 
 
+def minmax(a, axis=None, out=None, keepdims=False, initial=None, where=True):
+    """
+    Return the minimum and maximum of an array or along an axis.
+
+    This is equivalent to ``(dpnp.min(a, ...), dpnp.max(a, ...))`` but computes
+    both the minimum and the maximum in a single pass over `a`.
+
+    For full documentation refer to :obj:`numpy.minmax`.
+
+    Parameters
+    ----------
+    a : {dpnp.ndarray, usm_ndarray}
+        Input array.
+    axis : {None, int or tuple of ints}, optional
+        Axis or axes along which to operate. By default, flattened input is
+        used. If this is a tuple of integers, the minimum and the maximum are
+        selected over multiple axes, instead of a single axis or all the axes
+        as before.
+
+        Default: ``None``.
+    out : {None, tuple of None, dpnp.ndarray or usm_ndarray}, optional
+        A tuple ``(min, max)`` of two arrays in which to place the result.
+        Each array must be of the same shape and buffer length as the expected
+        output. A ``None`` entry means that a new array is allocated for that
+        result.
+
+        Default: ``None``.
+    keepdims : {None, bool}, optional
+        If this is set to ``True``, the axes which are reduced are left in the
+        result as dimensions with size one. With this option, the result will
+        broadcast correctly against the input array.
+
+        Default: ``False``.
+
+    Returns
+    -------
+    out : tuple of dpnp.ndarray
+        A tuple ``(min, max)`` holding the minimum and maximum of `a`. If
+        `axis` is ``None``, the results are zero-dimensional arrays. If `axis`
+        is an integer, the results are arrays of dimension ``a.ndim - 1``. If
+        `axis` is a tuple, the results are arrays of dimension
+        ``a.ndim - len(axis)``.
+
+    Limitations
+    -----------
+    Parameters `where`, and `initial` are only supported with their default
+    values. Otherwise ``NotImplementedError`` exception will be raised.
+
+    If an array in `out` has a data type other than that of `a`, the result is
+    computed in the data type of `a` and then cast, while NumPy casts `a` to
+    the data type of `out` first. The results differ when the cast does not
+    preserve the order of values, e.g. for a narrowing integer cast.
+
+    The arrays in `out` must not overlap in memory. Interleaved views, such as
+    ``(r[:, 0], r[:, 1])``, are rejected as well, while NumPy accepts them.
+
+    See Also
+    --------
+    :obj:`dpnp.min` : Return the minimum of an array.
+    :obj:`dpnp.max` : Return the maximum of an array.
+    :obj:`dpnp.ptp` : Range of values (maximum - minimum) along an axis.
+    :obj:`dpnp.nanmin` : The minimum value of an array along a given axis,
+                         ignores NaNs.
+    :obj:`dpnp.nanmax` : The maximum value of an array along a given axis,
+                         ignores NaNs.
+
+    Notes
+    -----
+    NaN values are propagated, that is if at least one item is NaN, the
+    corresponding output value will be NaN as well. To ignore NaN values
+    use :obj:`dpnp.nanmin` and :obj:`dpnp.nanmax`.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> a = np.arange(4).reshape((2, 2))
+    >>> a
+    array([[0, 1],
+           [2, 3]])
+    >>> np.minmax(a)            # min and max of the flattened array
+    (array(0), array(3))
+
+    >>> np.minmax(a, axis=0)    # along the first axis
+    (array([0, 1]), array([2, 3]))
+    >>> np.minmax(a, axis=1)    # along the second axis
+    (array([0, 2]), array([1, 3]))
+
+    >>> b = np.arange(5, dtype=float)
+    >>> b[2] = np.nan
+    >>> np.minmax(b)
+    (array(nan), array(nan))
+
+    """
+
+    dpnp.check_limitations(initial=initial, where=where)
+    usm_a = dpnp.get_usm_ndarray(a)
+
+    if out is None:
+        out = (None, None)
+    elif not isinstance(out, tuple):
+        raise TypeError("'out' must be a tuple of arrays")
+    elif len(out) != 2:
+        raise ValueError("'out' tuple must have exactly two entries")
+
+    if all(res is not None for res in out):
+        dpnp.check_supported_arrays_type(*out)
+        # pylint: disable=protected-access
+        if dti._array_overlap(*map(dpnp.get_usm_ndarray, out)):
+            raise ValueError("Output arrays cannot overlap")
+
+    usm_out = []
+    for res in out:
+        if res is not None:
+            dpnp.check_supported_arrays_type(res)
+            # the result is computed in the data type of `a`
+            if res.dtype != a.dtype:
+                res = dpnp.astype(res, a.dtype, copy=False)
+            res = dpnp.get_usm_ndarray(res)
+        usm_out.append(res)
+
+    usm_res = dpt_minmax(
+        usm_a, axis=axis, keepdims=keepdims, out=tuple(usm_out)
+    )
+    return tuple(
+        dpnp.get_result_array(x, res, casting="unsafe")
+        for x, res in zip(usm_res, out)
+    )
+
+
 def ptp(
     a,
     /,

dpnp/tensor/CMakeLists.txt

@@ -194,6 +194,7 @@ set(_reduction_sources
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/reductions/logsumexp.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/reductions/max.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/reductions/min.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/reductions/minmax.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/reductions/prod.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/reductions/reduce_hypot.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/reductions/sum.cpp

dpnp/tensor/_reduction.py

@@ -138,6 +138,139 @@ def _comparison_over_axis(x, axis, keepdims, out, _reduction_fn):
     return out
 
 
+def _minmax(x, /, *, axis=None, keepdims=False, out=None):
+    """
+    Calculates the minimum and the maximum of the input array ``x`` in a
+    single pass.
+
+    Args:
+        x (usm_ndarray):
+            input array.
+        axis (Optional[int, Tuple[int, ...]]):
+            axis or axes along which the extrema must be computed.
+            If ``None``, they are computed over the entire array.
+            Default: ``None``.
+        keepdims (Optional[bool]):
+            if ``True``, the reduced axes (dimensions) are included in the
+            results as singleton dimensions. Default: ``False``.
+        out (Optional[Tuple[Optional[usm_ndarray], Optional[usm_ndarray]]]):
+            the arrays into which the minima and the maxima are written.
+            Each array must have the expected shape and the data type of
+            ``x``. A ``None`` entry means that a new array is returned.
+            Default: ``None``.
+
+    Returns:
+        Tuple[usm_ndarray, usm_ndarray]:
+            the minima and the maxima of ``x``, with its data type.
+    """
+    if not isinstance(x, dpt.usm_ndarray):
+        raise TypeError(f"Expected dpnp.tensor.usm_ndarray, got {type(x)}")
+
+    nd = x.ndim
+    if axis is None:
+        axis = tuple(range(nd))
+        perm = list(axis)
+        x_tmp = x
+    else:
+        if not isinstance(axis, (tuple, list)):
+            axis = (axis,)
+        axis = normalize_axis_tuple(axis, nd, "axis")
+        perm = [i for i in range(nd) if i not in axis] + list(axis)
+        x_tmp = dpt.permute_dims(x, perm)
+    red_nd = len(axis)
+    if any([x_tmp.shape[i] == 0 for i in range(-red_nd, 0)]):
+        raise ValueError("reduction cannot be performed over zero-size axes")
+    res_shape = x_tmp.shape[: nd - red_nd]
+    exec_q = x.sycl_queue
+    res_dt = x.dtype
+    res_usm_type = x.usm_type
+
+    if keepdims:
+        final_res_shape = tuple(
+            x.shape[d] if d not in axis else 1 for d in range(nd)
+        )
+    else:
+        final_res_shape = res_shape
+
+    if out is None:
+        out = (None, None)
+    orig_out = list(out)
+    out = list(out)
+    for i, res in enumerate(out):
+        if res is None:
+            out[i] = dpt.empty(
+                res_shape,
+                dtype=res_dt,
+                usm_type=res_usm_type,
+                sycl_queue=exec_q,
+            )
+            continue
+
+        if not isinstance(res, dpt.usm_ndarray):
+            raise TypeError(
+                f"output array must be of usm_ndarray type, got {type(res)}"
+            )
+        if not res.flags.writable:
+            raise ValueError("provided `out` array is read-only")
+        if not res.shape == final_res_shape:
+            raise ValueError(
+                "The shape of input and output arrays are inconsistent. "
+                f"Expected output shape is {final_res_shape}, got {res.shape}"
+            )
+        if res_dt != res.dtype:
+            raise ValueError(
+                f"Output array of type {res_dt} is needed, got {res.dtype}"
+            )
+        if dpt.get_execution_queue((exec_q, res.sycl_queue)) is None:
+            raise dpt.ExecutionPlacementError(
+                "Input and output allocation queues are not compatible"
+            )
+        if keepdims:
+            res = dpt.squeeze(res, axis=axis)
+            orig_out[i] = res
+        if ti._array_overlap(x, res):
+            res = dpt.empty_like(res)
+        out[i] = res
+
+    _manager = SequentialOrderManager[exec_q]
+    dep_evs = _manager.submitted_events
+    if red_nd == 0:
+        for res in out:
+            ht_e_cpy, cpy_e = ti._copy_usm_ndarray_into_usm_ndarray(
+                src=x_tmp, dst=res, sycl_queue=exec_q, depends=dep_evs
+            )
+            _manager.add_event_pair(ht_e_cpy, cpy_e)
+        red_evs = _manager.submitted_events
+    else:
+        hev, red_ev = tri._minmax_over_axis(
+            src=x_tmp,
+            trailing_dims_to_reduce=red_nd,
+            dst_min=out[0],
+            dst_max=out[1],
+            sycl_queue=exec_q,
+            depends=dep_evs,
+        )
+        _manager.add_event_pair(hev, red_ev)
+        red_evs = [red_ev]
+
+    for i, res in enumerate(out):
+        if not (orig_out[i] is None or orig_out[i] is res):
+            ht_e_cpy2, cpy2_e = ti._copy_usm_ndarray_into_usm_ndarray(
+                src=res, dst=orig_out[i], sycl_queue=exec_q, depends=red_evs
+            )
+            _manager.add_event_pair(ht_e_cpy2, cpy2_e)
+            out[i] = orig_out[i]
+
+    if keepdims:
+        res_shape = res_shape + (1,) * red_nd
+        inv_perm = sorted(range(nd), key=lambda d: perm[d])
+        out = [
+            dpt.permute_dims(dpt.reshape(res, res_shape), inv_perm)
+            for res in out
+        ]
+    return tuple(out)
+
+
 def _reduction_over_axis(
     x,
     axis,