Function argmax

tensorrt_llm/functional.py:3303–3347

Add an argmax operation. As explained in the ONNX documentation, https://github.com/onnx/onnx/blob/main/docs/Operators.md#argmax that function creates a layer computing the indices of the max elements of the input tensor's element along the provided dim. The resulting ten

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

def argmax(input: Tensor, dim: int, keepdim: bool = False) -> Tensor:
    '''
    Add an argmax operation.

    As explained in the ONNX documentation,

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

    that function creates a layer computing the indices of the max elements of
    the input tensor's element along the provided dim. The resulting tensor
    has the same rank as the input if keepdims is True. If keepdims is False,
    then the resulting tensor has the reduced dimension pruned.

    Parameters:
        input : Tensor
            The input tensor.

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

        keepdim : bool
            Do we keep the dimension along which the reduction is performed?
            Yes, if set to True, no otherwise.

    Returns:
        The tensor produced by this argmax operation.
    '''
    dim = dim_resolve_negative(dim, input.ndim())
    axes = dim_to_trt_axes(dim)

    layer = default_trtnet().add_topk(input.trt_tensor, trt.TopKOperation.MAX,
                                      1, axes)
    output = layer.get_output(1)

    if keepdim:
        return _create_tensor(output, layer)

    output = _create_tensor(output, layer)
    a = list(range(input.ndim()))
    for d in dim:
        a.pop(d)
    indices = constant(int32_array(a))
    output_shape = shape(output)
    new_shape = gather(output_shape, 0, indices)
    return view(output, new_shape)


def gelu(x: Tensor) -> Tensor: