Function cumreduction

dask/array/reductions.py:1157–1278

Generic function for cumulative reduction Parameters ---------- func: callable Cumulative function like np.cumsum or np.cumprod binop: callable Associated binary operator like ``np.cumsum->add`` or ``np.cumprod->mul`` ident: Number Associated identity lik

(
    func,
    binop,
    ident,
    x,
    axis=None,
    dtype=None,
    out=None,
    method="sequential",
    preop=None,
)

def cumreduction(
    func,
    binop,
    ident,
    x,
    axis=None,
    dtype=None,
    out=None,
    method="sequential",
    preop=None,
):
    """Generic function for cumulative reduction

    Parameters
    ----------
    func: callable
        Cumulative function like np.cumsum or np.cumprod
    binop: callable
        Associated binary operator like ``np.cumsum->add`` or ``np.cumprod->mul``
    ident: Number
        Associated identity like ``np.cumsum->0`` or ``np.cumprod->1``
    x: dask Array
    axis: int
    dtype: dtype
    method : {'sequential', 'blelloch'}, optional
        Choose which method to use to perform the cumsum.  Default is 'sequential'.

        * 'sequential' performs the scan of each prior block before the current block.
        * 'blelloch' is a work-efficient parallel scan.  It exposes parallelism by first
          calling ``preop`` on each block and combines the values via a binary tree.
          This method may be faster or more memory efficient depending on workload,
          scheduler, and hardware.  More benchmarking is necessary.
    preop: callable, optional
        Function used by 'blelloch' method,
        like ``np.cumsum->np.sum`` or ``np.cumprod->np.prod``

    Returns
    -------
    dask array

    See also
    --------
    cumsum
    cumprod
    """
    if method == "blelloch":
        if preop is None:
            raise TypeError(
                'cumreduction with "blelloch" method required `preop=` argument'
            )
        return prefixscan_blelloch(func, preop, binop, x, axis, dtype, out=out)
    elif method != "sequential":
        raise ValueError(
            f'Invalid method for cumreduction.  Expected "sequential" or "blelloch".  Got: {method!r}'
        )

    if axis is None:
        if x.ndim > 1: