Function cumprod

python/tvm/topi/gpu/scan.py:729–776

Numpy style cumprod op. Return the cumulative product of the elements along a given axis. Parameters ---------- data : tvm.te.Tensor The input data to the operator. axis : int, optional Axis along which the cumulative product is computed. The default (None) is to co

(
    data: tvm.te.Tensor,
    axis: int | None = None,
    dtype: int | None = None,
    exclusive: bool | None = None,
    workspace: tvm.te.Tensor | None = None,
)

def cumprod(
    data: tvm.te.Tensor,
    axis: int | None = None,
    dtype: int | None = None,
    exclusive: bool | None = None,
    workspace: tvm.te.Tensor | None = None,
):
    """Numpy style cumprod op. Return the cumulative product of the elements along a given axis.

    Parameters
    ----------
    data : tvm.te.Tensor
        The input data to the operator.

    axis : int, optional
        Axis along which the cumulative product is computed. The default (None) is to compute
        the cumproduct over the flattened array.

    dtype : string, optional
        Type of the returned array and of the accumulator in which the elements are multiplied.
        If dtype is not specified, it defaults to the dtype of data.

    exclusive : bool, optional
        If True, will return exclusive product in which the first element is not
        included. In other terms, if True, the j-th output element would be
        the product of the first (j-1) elements. Otherwise, it would be the product of
        the first j elements.

    workspace: Optional[tvm.te.Tensor]
        A buffer to store intermediate results if thrust is enabled. The size of the workspace
        should be sufficiently large, this can be obtained by overestimation or memory usage
        profiling. If None, it will fallback to use thrust internal memory allocation.

    Returns
    -------
    result : tvm.te.Tensor
        The result has the same size as data, and the same shape as data if axis is not None.
        If axis is None, the result is a 1-d array.
    """
    return scanop(
        data=data,
        binop=tvm.tirx.generic.multiply,
        identity_value=1,
        axis=axis,
        dtype=dtype,
        exclusive=exclusive,
        workspace=workspace,
    )