Function split

tensorrt_llm/functional.py:3748–3821

Add an operation that splits a tensor into sub-tensors. This operation creates a list of tensors that are obtained from the input tensor by slicing it along the dimension 'dim'. If 'split_size_or_sections' is an integer, the tensor is split into 'input.shape[dim] / split_size_o

(tensor: Tensor,
          split_size_or_sections: Union[int, Sequence[int]],
          dim: int = 0)

def split(tensor: Tensor,
          split_size_or_sections: Union[int, Sequence[int]],
          dim: int = 0) -> Sequence[Tensor]:
    '''
    Add an operation that splits a tensor into sub-tensors.

    This operation creates a list of tensors that are obtained from the input
    tensor by slicing it along the dimension 'dim'. If 'split_size_or_sections'
    is an integer, the tensor is split into 'input.shape[dim] /
    split_size_or_sections' slices. If 'split_size_or_sections' is a list of
    sizes, the tensor is split into 'len(split_size_or_sections)' slices and
    the size of the ith slice is given by 'split_size_or_sections[i]'.

    There are several constraints with the current implementation:

        - The input tensor must be static (no dynamic dimension),
        - If 'split_size_or_sections' is an integer, the number of elements in
          the 'dim' dimension of the input must be a multiple of
          'split_size_or_sections': 'input.shape[dim] % split_size_or_sections == 0'.
        - If 'split_size_or_sections' is a sequence, the sum of the elements in
          'split_size_or_sections' must be equal to the size in the dimension
          'dim': 'input.shape[dim] == sum(ii for ii in split_size_or_sections)'.

    That operation is implemented using a 'slice' operation for each output
    slice.

    Parameters:
        tensor : Tensor
            The input tensor to slice.

        split_size_or_sections : Union[int, Sequence[int]]
            If it is an integer, it encodes the size of each slice. Otherwise,
            if it is a sequence, it is the size of each slice.

        dim : int
            The dimension of the tensor to slice.

    Returns:
        The list of tensors produced by the different operations.
    '''
    assert not tensor.is_dynamic(dim)

    ndim = tensor.ndim()
    if dim < 0:
        dim += ndim
    dim_value = tensor.size()[dim]
    starts = [constant(dims_array([0])) for _ in range(ndim)]
    sizes = [shape(tensor, i) for i in range(ndim)]

    if isinstance(split_size_or_sections, int):
        # TODO: support non-divisible cases
        assert dim_value % split_size_or_sections == 0
        num_sections = dim_value // split_size_or_sections
        sizes[dim] = constant(dims_array([split_size_or_sections]))

        outputs = []
        for i in range(num_sections):
            starts[dim] = constant(dims_array([split_size_or_sections * i]))