Function reduce

tensorrt_llm/functional.py:3160–3194

Add an reduction operation to do along a dimension. It is implemented using the IReduceLayer from TensorRT. Parameters: input : Tensor The input tensor. op : trt.ReduceOperation The reduction operation to perform. Options: SUM, PROD

(input: Tensor,
           op: trt.ReduceOperation,
           dim: Union[int, Tuple[int]],
           keepdim: bool = False)

def reduce(input: Tensor,
           op: trt.ReduceOperation,
           dim: Union[int, Tuple[int]],
           keepdim: bool = False) -> Tensor:
    '''
    Add an reduction operation to do along a dimension.

    It is implemented using the IReduceLayer from TensorRT.

    Parameters:
        input : Tensor
            The input tensor.

        op : trt.ReduceOperation
            The reduction operation to perform.
            Options: SUM, PROD, MAX, MIN, AVG

        dim : int
            The dimension along which the reduction is performed.

        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.
    '''
    dim = dim_resolve_negative(dim, input.ndim())
    axes = dim_to_trt_axes(dim)

    layer = default_trtnet().add_reduce(input.trt_tensor,
                                        op,
                                        axes,
                                        keep_dims=keepdim)
    return _create_tensor(layer.get_output(0), layer)


prod = partial(reduce, op=trt.ReduceOperation.PROD)