Function cumsum

tensorrt_llm/functional.py:2411–2536

Add an operation to calculate inclusive cumulative sum of elements of a tensor in a given dimension. Given an input tensor, that function creates an operation that calculates inclusive cumulative sum of elements in the dimension 'dim' to create a new tensor. The output tensor h

(input: Tensor, dim: int, prefer_plugin: bool = True)

def cumsum(input: Tensor, dim: int, prefer_plugin: bool = True) -> Tensor:
    '''
    Add an operation to calculate inclusive cumulative sum of elements of
    a tensor in a given dimension.

    Given an input tensor, that function creates an operation that calculates
    inclusive cumulative sum of elements in the dimension 'dim' to create
    a new tensor. The output tensor has the same shape as the input tensor.

    The input tensor must have rank >= 1. The 'dim' must be valid, and negative
    value is supported.

    For example, on input=[[4, 2, 5], [2, 1, 2], [4, 7, 1]], which has a shape
    [3, 3],

        cumsum(input, 0)

    will produce [[4, 2, 5], [6, 3, 7], [10, 10, 8]].

        cumsum(input, 1)

    will produce [[4, 6, 11], [2, 3, 5], [4, 11, 12]].

    That operation is implemented by TensorRT ILoopLayer.

    Parameters:
        input : Tensor
            The input tensor to calculate the inclusive cumulative sum.

        dim : int
            The dimension to calculate the inclusive cumulative sum. Negative
            value is supported.

        prefer_plugin : bool
            Whether to use the cumsumLastDim plugin if dim is last dim.

    Returns:
        The tensor containing the inclusive cumulative sum of input.
    '''
    assert input.rank() >= 1, "input should have rank >= 1"
    assert dim < input.rank() and dim >= -input.rank(
    ), f"dim should be in [{-input.rank()}, {input.rank()}) when input have rank {input.rank()}"

    dim = dim_resolve_negative(dim, input.ndim())[0]

    if dim == input.ndim() - 1:
        if prefer_plugin:
            last_dim = input.size(-1)
            if last_dim == -1:  # dynamic?
                last_dim = shape(input, -1)
            old_shape = shape(input)
            if input.ndim() == 1:
                input_2d = unsqueeze(
                    input, 0)  # special handling of rank-1 dynamic tensor
            elif input.ndim() != 2:
                input_2d = input.view(concat([-1, last_dim]),
                                      zero_is_placeholder=False)
            else: