Function index_select

tensorrt_llm/functional.py:2216–2275

Add an operation to select slices of elements from a tensor. Given an input tensor, that function creates an operation that selects the slices of elements in the dimension 'dim' at the indices listed in 'index' to create a new tensor. The output tensor has the same rank as the inp

(input: Tensor, dim: int, index: Tensor)

def index_select(input: Tensor, dim: int, index: Tensor) -> Tensor:
    '''
    Add an operation to select slices of elements from a tensor.

    Given an input tensor, that function creates an operation that selects the
    slices of elements in the dimension 'dim' at the indices listed in 'index'
    to create a new tensor.  The output tensor has the same rank as the input
    tensor.

    The 'index' is a tensor of rank 1.

    For example, on input=[[4, 2, 5], [2, 1, 2], [4, 7, 1]], which has a shape
    [3, 3],

        index_select(input, 0, [0, 1])

    will create a tensor of shape [2, 3] that contains the [[4, 2, 5], [2, 1, 2]].

    Regarding the shape of the output tensor, the dimension 'dim' has the same
    size as the 'index' tensor. It means that for a input tensor of shape [4, 2, 6, 3],

        index_select(input, 2, [1, 4])

    will select the 2nd and 5th slices (index == 1 or 4) from the 3rd dimension
    (dim == 2) and return a tensor of shape [4, 2, 2, 3] (i.e. the 3rd
    dimension is shrunk to 2).

    Note that this operation can also be used to expand a tensor in the 'dim'
    dimension, for example, on input [[0, 1], [2, 3]],

        index_select(input, 1, [0, 0, 0])

    will produce a tensor of shape [2, 3] containing [[0, 0, 0], [2, 2, 2]].

    That operation maps to the TensorRT IGatherLayer.

    Parameters:
        input : Tensor
            The input tensor to select from.

        dim : int
            The dimension to select from.

        index : Tensor
            The indices of the slices in the 'dim' dimension to select.

    Returns:
        The tensor containing the selected slices.
    '''
    assert index.rank() == 1, f"index should have rank 1, got {index.rank()}"

    new_shape = []
    for i in range(input.rank()):
        if i != dim:
            new_shape.append(shape(input, i))
        else:
            new_shape.append(shape(index, 0))