Function sum

tensorrt_llm/functional.py:3253–3275

Add an operation to compute the sum along a dimension. Computes the sum along the dimension 'dim' of the input tensor. It is implemented using the IReduceLayer from TensorRT. Parameters: input : Tensor The input tensor. dim : int The dimen

(input: Tensor, dim: int, keepdim: bool = False)

def sum(input: Tensor, dim: int, keepdim: bool = False) -> Tensor:
    '''
    Add an operation to compute the sum along a dimension.

    Computes the sum along the dimension 'dim' of the input tensor.

    It is implemented using the IReduceLayer from TensorRT.

    Parameters:
        input : Tensor
            The input tensor.

        dim : int
            The dimension along which the mean is computed.

        keepdim : bool
            Is the dimension kept in the reduced tensor? When True the
            dimension is kept, it is removed from the shape otherwise.

    Returns:
        The tensor produced by this reduction operation.
    '''
    return reduce(input, op=trt.ReduceOperation.SUM, dim=dim, keepdim=keepdim)


def identity(input: Tensor) -> Tensor: