flopscope.

flopscope.numpy.apply_over_axes

fnp.apply_over_axes(func, a, axes)[flopscope source][numpy source]

Apply a function repeatedly over multiple axes.

Adapted from NumPy docs np.apply_over_axes

Areacore
Typecustom
NumPy Refnp.apply_over_axes
Cost
per-operation
Flopscope Context

Apply function over multiple axes. Cost: numel(output).

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.

Parameters

func:function

This function must take two arguments, func(a, axis).

a:array_like

Input array.

axes:array_like

Axes over which func is applied; the elements must be integers.

Returns

apply_over_axis: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

Notes

This function is equivalent to tuple axis arguments to reorderable ufuncs with keepdims=True. Tuple axis arguments to ufuncs have been available since version 1.7.0.

Examples

>>> import flopscope.numpy as fnp
>>> a = flops.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:

>>> flops.apply_over_axes(flops.sum, a, [0,2])
array([[[ 60],
        [ 92],
        [124]]])

Tuple axis arguments to ufuncs are equivalent:

>>> flops.sum(a, axis=(0,2), keepdims=True)
array([[[ 60],
        [ 92],
        [124]]])