Function scanop

python/tvm/topi/gpu/scan.py:604–676

Cumulative binary operator (scan) with similar axis behavior as np.cumsum and np.cumprod. See cumprod and cumsum for an example of use. E.g. if * is your binary operator and the input tensor is [1, 2, 3, 4] the output may be [1, 1 * 2, 1 * 2 * 3, 1 * 2 * 3 * 4] Parameters ----

(
    data: tvm.te.Tensor,
    binop: Callable[["tvm.Expr", "tvm.Expr"], "tvm.Expr"],
    identity_value: float | int,
    axis: int | None = None,
    dtype: str | None = None,
    exclusive: bool | None = None,
    workspace: tvm.te.Tensor | None = None,
)

def scanop(
    data: tvm.te.Tensor,
    binop: Callable[["tvm.Expr", "tvm.Expr"], "tvm.Expr"],
    identity_value: float | int,
    axis: int | None = None,
    dtype: str | None = None,
    exclusive: bool | None = None,
    workspace: tvm.te.Tensor | None = None,
) -> tvm.te.Tensor:
    """Cumulative binary operator (scan) with similar axis behavior as np.cumsum and np.cumprod.

    See cumprod and cumsum for an example of use.

    E.g. if * is your binary operator and the input tensor is [1, 2, 3, 4] the output may be
    [1, 1 * 2, 1 * 2 * 3, 1 * 2 * 3 * 4]

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

    binop: Callable (tvm.Expr, tvm.Expr) -> tvm.Expr
        A binary operator which should be associative and commutative. E.g. if * is your
        operator then a * (b * c) = (a * b) * c and a * b = b * a

    identity_value: int or float
        A value for the binary operation which provides the identity property. E.g. if * is
        your operator and i is the identity_value then a * i = a for all a in the domain of
        your operation.

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

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

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

    workspace: Optional[tvm.te.Tensor]

    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.
    """
    if axis is None:
        axis = 0
        data = reshape(data, (prod(data.shape),))
    axis = get_const_int(axis)
    if exclusive is not None and exclusive:
        return exclusive_scan(