Function flatten

tensorrt_llm/functional.py:1793–1828

Flattens input by reshaping it into a one-dimensional tensor. If start_dim or end_dim are passed, only dimensions starting with start_dim and ending with end_dim are flattened. The order of elements in input is unchanged. Parameters: input : Tensor The input te

(input: Tensor, start_dim: int = 0, end_dim: int = -1)

def flatten(input: Tensor, start_dim: int = 0, end_dim: int = -1):
    '''
    Flattens input by reshaping it into a one-dimensional tensor.

    If start_dim or end_dim are passed, only dimensions starting with start_dim and
    ending with end_dim are flattened. The order of elements in input is unchanged.

    Parameters:
        input : Tensor
            The input tensor to flatten.

        start_dim : int
            The first dim to flatten.

        end_dim : int
            The last dim to flatten.

    Returns:
        The tensor produced by the flatten layer.

    '''
    shape = input.shape
    ndim = input.ndim()
    if start_dim < 0: start_dim += ndim
    if end_dim < 0: end_dim += ndim
    new_shape = list()
    for i in range(start_dim):
        new_shape.append(shape[i])
    if end_dim - start_dim >= 0:
        flat_dim = 1
        for i in range(start_dim, end_dim + 1):
            flat_dim *= shape[i]
        new_shape.append(flat_dim)
    for i in range(end_dim + 1, ndim):
        new_shape.append(shape[i])
    return view(input, new_shape)


def expand_dims(input: Tensor,