Function max

tensorrt_llm/functional.py:3228–3250

Add an operation to compute the max along a dimension. Computes the max 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 max(input: Tensor, dim: int, keepdim: bool = False) -> Tensor:
    '''
    Add an operation to compute the max along a dimension.

    Computes the max 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.MAX, dim=dim, keepdim=keepdim)


def sum(input: Tensor, dim: int, keepdim: bool = False) -> Tensor: