Function permute

tensorrt_llm/functional.py:1685–1712

Add an operation to permute the dimensions of a tensor. The dimensions of the input tensor are permuted according to the sequence of dimensions in 'dims'. That operation maps to tensorrt.IShuffleLayer where the second transposition is described by the indices in 'dims'. Given

(input: Tensor, dims: Sequence[int])

def permute(input: Tensor, dims: Sequence[int]) -> Tensor:
    '''
    Add an operation to permute the dimensions of a tensor.

    The dimensions of the input tensor are permuted according to the sequence
    of dimensions in 'dims'. That operation maps to tensorrt.IShuffleLayer where
    the second transposition is described by the indices in 'dims'.

    Given a tensor of rank N, the result of the permutation is a tensor of rank
    N in which the i-th input dimension maps to the dims[i]-th dimension.

    For example, permute(input, [1, 0]) will transpose a 2D tensor by permuting
    the rows and columns.

    Parameters:
        input : Tensor
            The input tensor to permute.

        dims : Sequence[int]
            The description of the permutation.

    Returns:
        The tensor produced by the permutation layer.
    '''
    dims = dim_resolve_negative(tuple(dims), input.ndim())
    layer = default_trtnet().add_shuffle(input.trt_tensor)
    layer.second_transpose = dims
    return _create_tensor(layer.get_output(0), layer)


def transpose(input: Tensor, dim0: int, dim1: int) -> Tensor: