Merge branch 'master' into dependabot/github_actions/github/codeql-action-3.27.4

Author: vtavana

dpnp/dpnp_iface_functional.py

@@ -38,11 +38,14 @@
 
 
 import numpy
-from dpctl.tensor._numpy_helper import normalize_axis_index
+from dpctl.tensor._numpy_helper import (
+    normalize_axis_index,
+    normalize_axis_tuple,
+)
 
 import dpnp
 
-__all__ = ["apply_along_axis"]
+__all__ = ["apply_along_axis", "apply_over_axes"]
 
 
 def apply_along_axis(func1d, axis, arr, *args, **kwargs):
@@ -185,3 +188,83 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs):
         buff = dpnp.moveaxis(buff, -1, axis)
 
     return buff
+
+
+def apply_over_axes(func, a, axes):
+    """
+    Apply a function repeatedly over multiple axes.
+
+    `func` is called as ``res = func(a, axis)``, where `axis` is the first
+    element of `axes`. The result `res` of the function call must have
+    either the same dimensions as `a` or one less dimension. If `res`
+    has one less dimension than `a`, a dimension is inserted before
+    `axis`. The call to `func` is then repeated for each axis in `axes`,
+    with `res` as the first argument.
+
+    For full documentation refer to :obj:`numpy.apply_over_axes`.
+
+    Parameters
+    ----------
+    func : function
+         This function must take two arguments, ``func(a, axis)``.
+    a : {dpnp.ndarray, usm_ndarray}
+        Input array.
+    axes : {int, sequence of ints}
+        Axes over which `func` is applied.
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        The output array. The number of dimensions is the same as `a`,
+        but the shape can be different. This depends on whether `func`
+        changes the shape of its output with respect to its input.
+
+    See Also
+    --------
+    :obj:`dpnp.apply_along_axis` : Apply a function to 1-D slices of an array
+                                   along the given axis.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> a = np.arange(24).reshape(2, 3, 4)
+    >>> a
+    array([[[ 0,  1,  2,  3],
+            [ 4,  5,  6,  7],
+            [ 8,  9, 10, 11]],
+           [[12, 13, 14, 15],
+            [16, 17, 18, 19],
+            [20, 21, 22, 23]]])
+
+    Sum over axes 0 and 2. The result has same number of dimensions
+    as the original array:
+
+    >>> np.apply_over_axes(np.sum, a, [0, 2])
+    array([[[ 60],
+            [ 92],
+            [124]]])
+
+    Tuple axis arguments to ufuncs are equivalent:
+
+    >>> np.sum(a, axis=(0, 2), keepdims=True)
+    array([[[ 60],
+            [ 92],
+            [124]]])
+
+    """
+
+    dpnp.check_supported_arrays_type(a)
+    if isinstance(axes, int):
+        axes = (axes,)
+    axes = normalize_axis_tuple(axes, a.ndim)
+
+    for axis in axes:
+        res = func(a, axis)
+        if res.ndim != a.ndim:
+            res = dpnp.expand_dims(res, axis)
+            if res.ndim != a.ndim:
+                raise ValueError(
+                    "function is not returning an array of the correct shape"
+                )
+        a = res
+    return res

dpnp/dpnp_iface_manipulation.py

@@ -3848,13 +3848,14 @@ def transpose(a, axes=None):
     ----------
     a : {dpnp.ndarray, usm_ndarray}
         Input array.
-    axes : None, tuple or list of ints, optional
+    axes : {None, tuple or list of ints}, optional
         If specified, it must be a tuple or list which contains a permutation
         of [0, 1, ..., N-1] where N is the number of axes of `a`.
         The `i`'th axis of the returned array will correspond to the axis
         numbered ``axes[i]`` of the input. If not specified or ``None``,
         defaults to ``range(a.ndim)[::-1]``, which reverses the order of
         the axes.
+        Default: ``None``.
 
     Returns
     -------
@@ -3894,18 +3895,13 @@ def transpose(a, axes=None):
 
     """
 
-    if isinstance(a, dpnp_array):
-        array = a
-    elif isinstance(a, dpt.usm_ndarray):
-        array = dpnp_array._create_from_usm_ndarray(a)
-    else:
-        raise TypeError(
-            f"An array must be any of supported type, but got {type(a)}"
-        )
+    dpnp.check_supported_arrays_type(a)
+    if isinstance(a, dpt.usm_ndarray):
+        a = dpnp_array._create_from_usm_ndarray(a)
 
     if axes is None:
-        return array.transpose()
-    return array.transpose(*axes)
+        return a.transpose()
+    return a.transpose(*axes)
 
 
 permute_dims = transpose  # permute_dims is an alias for transpose