Function unsqueeze

tensorrt_llm/functional.py:1930–1958

Add an operation to insert a singleton dimension to a tensor. That functions creates an operation that insert a singleton dimension (dimension of size 1) at position 'axis' in the output tensor. It works with negative values for the 'axis'. For example, for a tensor 'input' of

(input: Tensor, axis: int)

def unsqueeze(input: Tensor, axis: int):
    '''
    Add an operation to insert a singleton dimension to a tensor.

    That functions creates an operation that insert a singleton dimension
    (dimension of size 1) at position 'axis' in the output tensor. It works with
    negative values for the 'axis'.

    For example, for a tensor 'input' of shape [4, 4]:

        unsqueeze(input,  0) will produce an output of shape [1, 4, 4],
        unsqueeze(input,  1) will produce an output of shape [4, 1, 4],
        unsqueeze(input, -1) will produce an output of shape [4, 4, 1],
        unsqueeze(input, -2) will produce an output of shape [4, 1, 4],

    Parameters:
        input : Tensor
            The input tensor to expand with a singleton dimension.

        axis : int
            The index of the singleton dimension in the output tensor.

    Returns:
        The tensor produced by the layer.
    '''
    if axis < 0:
        axis = axis + input.ndim() + 1

    return expand_dims(input, axis)


def stack(inputs: Sequence[Tensor], dim: int = 0) -> Tensor: