Function topk

tensorrt_llm/functional.py:7320–7416

Add an topk operation. As explained in the ONNX documentation, https://github.com/onnx/onnx/blob/main/docs/Operators.md#topk NOTE: One distinction from the ONNX topk op, the output is always sorted with TensorRT layer. Retrieve the top-K largest elements along a spec

(input: Tensor,
         k: Union[Tensor, int],
         dim: int,
         largest: bool = True,
         prefer_plugin: bool = True)

def topk(input: Tensor,
         k: Union[Tensor, int],
         dim: int,
         largest: bool = True,
         prefer_plugin: bool = True) -> Tuple[Tensor, Tensor]:
    '''
    Add an topk operation.

    As explained in the ONNX documentation,

        https://github.com/onnx/onnx/blob/main/docs/Operators.md#topk

    NOTE: One distinction from the ONNX topk op, the output is always sorted
    with TensorRT layer.

    Retrieve the top-K largest elements along a specified axis.
    Given an input tensor of shape [a_1, a_2, ..., a_n, r]
    and integer argument k, return two outputs:
    Value tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] which contains the values of the top k elements along the specified axis
    Index tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] which contains the indices of the top k elements (original indices from the input tensor).

    Parameters:
        input : Tensor
            The input tensor.

        k : int
            A single positive value corresponding to the number of top elements to retrieve

        dim: int
            The dimension in which to compute the topk indices.

        largest: bool
            Controls whether to return largest or smallest elements

        prefer_plugin : bool
            Whether to use the topkLastDim plugin if dim is last dim and k is static.


    Returns:
        The tensors (values, indices) produced by this topk operation.
    '''
    dim = dim_resolve_negative(dim, input.ndim())[0]
    if prefer_plugin and dim == input.ndim() - 1 and not isinstance(k, Tensor):
        last_dim = input.size(-1)
        if last_dim == -1:  # dynamic?
            last_dim = shape(input, -1)
        # since we might need to flatten the input to 2d tensor,
        # we need to prepare the output shape
        out_shape = []
        for i in range(input.ndim() - 1):
            out_shape.append(shape(input, i))
        out_shape = concat(out_shape + [k])
        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)