Function transpose

tensorrt_llm/functional.py:1715–1742

Add an operation to transpose two dimensions of a tensor. That operation produces a tensor in which the dimensions 'dim0' and 'dim1' are permuted. The other dimensions, if the rank of the tensor is greater than 2, remain untouched. That function is a helper built on the 'funct

(input: Tensor, dim0: int, dim1: int)

def transpose(input: Tensor, dim0: int, dim1: int) -> Tensor:
    '''
    Add an operation to transpose two dimensions of a tensor.

    That operation produces a tensor in which the dimensions 'dim0' and 'dim1'
    are permuted. The other dimensions, if the rank of the tensor is greater
    than 2, remain untouched.

    That function is a helper built on the 'functional.permute' function.

    Parameters:
        input : Tensor
            The input tensor to transpose.

        dim0 : int
            The first dimension to transpose.

        dim1 : int
            The second dimension to transpose.

    Returns:
        The tensor produced by the permutation layer.
    '''
    permutation = list(range(input.ndim()))
    permutation[dim0] = dim1
    permutation[dim1] = dim0

    return permute(input, permutation)


def view(input: Tensor,