Function squeeze

tensorrt_llm/functional.py:1886–1927

Add an operation to remove singleton dimensions of a tensor. This functions creates an operation that removes singleton dimension (dimension of size 1) at positions 'dim' in the input tensor. It works with negative values for the 'dim'. For example, for a tensor 'input' of sha

(input: Tensor,
            dim: Optional[Union[int, Sequence[int]]] = None,
            zero_is_placeholder: bool = False)

# NOTE: Jointly added with Apple
def squeeze(input: Tensor,
            dim: Optional[Union[int, Sequence[int]]] = None,
            zero_is_placeholder: bool = False):
    '''
    Add an operation to remove singleton dimensions of a tensor.

    This functions creates an operation that removes singleton dimension
    (dimension of size 1) at positions 'dim' in the input tensor. It works with
    negative values for the 'dim'.

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

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

    Parameters:
        input : Tensor
            The input tensor for which the singleton dimensions will be removed.

        dim : Union[int, Sequence[int]]
            The index of the singleton dimensions in the input tensor.

    Returns:
        The tensor produced by the layer.
    '''
    if dim is None:
        dim = list(range(input.ndim()))
    if isinstance(dim, int):
        dim = (dim, )
    dim = dim_resolve_negative(dim, input.ndim())

    new_shape = []
    for i, s in enumerate(input.shape):
        if s == 1 and i in dim:
            continue
        new_shape.append(shape(input, i))

    new_shape = concat(new_shape) if len(new_shape) > 0 else []
    input = input.view(new_shape, zero_is_placeholder=zero_is_placeholder)
    return input


def unsqueeze(input: Tensor, axis: int):