Function repeat_interleave

tensorrt_llm/functional.py:6314–6338

Repeats elements of a tensor along an axis. Parameters: repeats : int The number of repetitions along axis specified. dim : int The dimension along which repetitions are performed. Returns: A tensor with the same shape as input except fo

(tensor: Tensor, repeats: int, dim: int)

def repeat_interleave(tensor: Tensor, repeats: int, dim: int) -> Tensor:
    '''
    Repeats elements of a tensor along an axis.

    Parameters:
        repeats : int
            The number of repetitions along axis specified.
        dim : int
            The dimension along which repetitions are performed.

    Returns:
        A tensor with the same shape as input except for repeated elements along specified dim.

    TODO: Allow repeats to be a list of integers and dim to be unspecified.
    '''
    expanded_tensor = expand_dims(tensor, dim + 1)
    tile_output_size = concat([
        repeats if i == (dim + 1) else shape(expanded_tensor, i)
        for i in range(expanded_tensor.ndim())
    ])
    tile = expand(expanded_tensor, tile_output_size)
    tile_reshape_size = [shape(tensor, i) for i in range(tensor.ndim())]
    tile_reshape_size[dim] = tile_reshape_size[dim] * repeats
    tensor = tile.view(concat(tile_reshape_size))
    return tensor


def meshgrid2d(x: Tensor, y: Tensor) -> Tuple[Tensor]: