Function softmax

tensorrt_llm/functional.py:2639–2667

Add an operation to compute softmax on a tensor. That operation computes the softmax on the input tensor in the dimension 'dim' if specified. Otherwise, it is applied on the last dimension. It inserts a ISoftmaxLayer to the TensorRT graph. Parameters: input : Tensor

(input: Tensor, dim: Optional[int] = None)

def softmax(input: Tensor, dim: Optional[int] = None) -> Tensor:
    '''
    Add an operation to compute softmax on a tensor.

    That operation computes the softmax on the input tensor in the dimension
    'dim' if specified. Otherwise, it is applied on the last dimension.

    It inserts a ISoftmaxLayer to the TensorRT graph.

    Parameters:
        input : Tensor
            The input tensor on which to apply softmax.

        dim : Optional[int]
            The dimension used to apply softmax.

    Returns:
        The output tensor of the softmax layer.
    '''
    if dim is None:
        dim = input.ndim() - 1
    if dim < 0:
        dim = input.ndim() + dim
    axes = dim_to_trt_axes(dim)

    layer = default_trtnet().add_softmax(input.trt_tensor)
    layer.axes = axes

    return _create_tensor(layer.get_output(0), layer)


def _lookup_plugin(input: Tensor, weight: Tensor, rank: int,