Function slice

tensorrt_llm/functional.py:1222–1324

Add an operation to extract a slice from a tensor. As described in the TensorRT documentation of the ISliceLayer, the slice layer has two variants: Static and dynamic. For static slicing, this function takes the starts and sizes values in the different dimensions to slice at l

(input: Tensor,
          starts: Union[Tensor, Sequence[int]],
          sizes: Union[Tensor, Sequence[int]],
          strides: Union[Tensor, Sequence[int]] = None,
          mode: trt.SampleMode = None,
          fill_value: Union[float, Tensor] = None)

# TODO: TensorRT uses sizes of the output dimensions.
# DL framework uses ends usually. Will change it to ends.
def slice(input: Tensor,
          starts: Union[Tensor, Sequence[int]],
          sizes: Union[Tensor, Sequence[int]],
          strides: Union[Tensor, Sequence[int]] = None,
          mode: trt.SampleMode = None,
          fill_value: Union[float, Tensor] = None) -> Tensor:
    '''
    Add an operation to extract a slice from a tensor.

    As described in the TensorRT documentation of the ISliceLayer, the slice
    layer has two variants: Static and dynamic.

    For static slicing, this function takes the starts and sizes values in the
    different dimensions to slice at layer creation time via a sequence of
    integers. For dynamic slicing, it accepts starts and sizes as
    tensorrt.ITensor`s.

    The slice layer selects for each dimension a start location from within the
    input tensor, and copies elements to the output tensor using a stride of 1
    across the input tensor. Start and size tensors must be 1-D int32 shape
    tensors if not specified as a sequence of integers.

    As an example, on input = [[0, 2, 4], [1, 3, 5]], the call to

        slice(input, start=[1, 0], size=[1, 2])

    will produce the tensor [[1, 3]] as output. The slice operator when
    executed by TensorRT will copy one row (because size[0] == 1) starting from
    the 2nd row (because start[0] == 1) and two columns (size[1] == 2) starting
    from the 1st column (because start[1] == 0).

    In pseudo-code the behavior of that operation can be described as follows
    for a 2D tensor (and easily be extended to more dimensions):

        output = Tensor(shape=sizes)
        for ii in range(sizes[0]):
            for jj in range(sizes[1]):
                output[ii][jj] = input[starts[0]+ii][starts[1]+jj]

    Note that it is common in deep-learning frameworks to use ranges
    [start:end] for similar operations. It can be emulated by setting the sizes
    argument such that in each dimension [start:start+size] == [start:end] i.e.
    size = end-start.

    TensorRT supports different slice modes but that function restricts that
    choice to `mode == tensorrt.SampleMode.STRICT_BOUNDS`.

    Parameters:
        input : Tensor
            The input tensor on which the slicing is performed.

        starts : Union[Tensor, Sequence[int]]
            The starting points, in the input tensor, and each dimension.

        sizes : Union[Tensor, Sequence[int]]
            The number of elements in each dimension of the sliced tensor (output).

        strides : Union[Tensor, Sequence[int]]