Function select

tensorrt_llm/functional.py:2160–2213

Add an operation to select a slice of elements from a tensor. Given an input tensor, that function creates an operation that selects the index-th slice of elements in the dimension 'dim' to create a new tensor. The output tensor has a shape in which the input dimension 'dim' is

(input: Tensor, dim: int, index: Union[Tensor, int])

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

    Given an input tensor, that function creates an operation that selects the
    index-th slice of elements in the dimension 'dim' to create a new tensor.
    The output tensor has a shape in which the input dimension 'dim' is
    removed.

    The 'index' can either be an integer or a 1D tensor containing a single
    element.

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

        select(input, 0, 1)

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

    Regarding the shape of the output tensor, the dimension 'dim' is removed.
    It means that for a tensor of shape [4, 2, 6, 3],

        select(input, 2, 4)

    will select the 5th slice (index == 4) from the 3rd dimension (dim == 2)
    and return a tensor of shape [4, 2, 3] (i.e. the 3rd dimension is removed).

    That operation maps to the TensorRT IGatherLayer.

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

        dim : int
            The dimension to select from.

        index : Union[Tensor, int]
            The index of the slice in the 'dim' dimension to select.

    Returns:
        The tensor containing the selected slice.
    '''
    if isinstance(index, int):
        index = constant(int32_array([index]))
    assert index.rank() == 1 and index.size(
        0) == 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))

    layer = default_trtnet().add_gather(input.trt_tensor, index.trt_tensor, dim)
    return _create_tensor(layer.get_output(0), layer).view(concat(new_shape))


def index_select(input: Tensor, dim: int, index: Tensor) -> Tensor: