Function repeat

tensorrt_llm/functional.py:6292–6311

Repeats the tensor along the specified dimensions. Parameters: input : Tensor The tensor to be repeated. sizes : Sequence[int] The number of times to repeat the tensor along each dimension. Returns: A tensor except for repeated input ten

(input: Tensor, sizes: Sequence[int])

def repeat(input: Tensor, sizes: Sequence[int]) -> Tensor:
    '''
    Repeats the tensor along the specified dimensions.

    Parameters:
        input : Tensor
            The tensor to be repeated.
        sizes : Sequence[int]
            The number of times to repeat the tensor along each dimension.

    Returns:
        A tensor except for repeated input tensors along specified dim.

    '''
    assert input.ndim() <= len(sizes), \
        "Number of dimensions of repeat dims can not be smaller than number of dimensions of tensor"
    repeated_tensor = input
    for k in range(-1, -len(sizes) - 1, -1):
        repeated_tensor = concat([repeated_tensor] * sizes[k], dim=k)
    return repeated_tensor


def repeat_interleave(tensor: Tensor, repeats: int, dim: int) -> Tensor: