Function expand_dims

tensorrt_llm/functional.py:1831–1882

Add an operation to expand the tensor shape with singleton dimensions. That function adds a tensorrt.IShuffleLayer to the network. Given an 'input' of rank N and a sequence of M dimensions, the output tensor produced by this operation (when executed by TensorRT) will have a rank of

(input: Tensor,
                dim: Union[int, Sequence[int]],
                shape_cast_dtype=None)

def expand_dims(input: Tensor,
                dim: Union[int, Sequence[int]],
                shape_cast_dtype=None) -> Tensor:
    '''
    Add an operation to expand the tensor shape with singleton dimensions.

    That function adds a tensorrt.IShuffleLayer to the network. Given an 'input'
    of rank N and a sequence of M dimensions, the output tensor produced by
    this operation (when executed by TensorRT) will have a rank of N+M. Singleton
    dimensions will be inserted at the different positions in 'dim'.

    The pseudo-code for that operation is:

        new_shape, ii = [], 0
        for jj in range(input.rank() + len(dim)):
            new_shape.append(1 if jj in dims else input.shape[ii++])

    For example, for a tensor of shape [3, 4, 1, 5]

        expand_dims(input, [0, 2])

    will produce a tensor of shape [1, 3, 1, 4, 1, 5].

    Parameters:
        input : Tensor
            The input tensor to expand.

        dim : Union[int, Sequence[int]]
            The positions in the output tensor where to insert singleton
            dimensions.

    Returns:
        The tensor produced by the shuffle layer.
    '''
    if isinstance(dim, int):
        dim = (dim, )

    out_ndim = len(dim) + input.ndim()

    input_shape = shape(input, cast_to_dtype=shape_cast_dtype)
    out_shapes = []
    j = 0
    for i in range(out_ndim):
        if i in dim:
            out_shapes.append(1)
        else:
            out_shapes.append(gather(input_shape, 0, j))
            j = j + 1

    out_shape = concat(out_shapes)

    return view(input, out_shape, zero_is_placeholder=False)


# NOTE: Jointly added with Apple