Function chunk

tensorrt_llm/functional.py:3824–3859

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 chunking it along the dimension 'dim'. It produces 'chunks' sub-tensors. That operation is only defined for static tensors (no dynamic di

(tensor: Tensor, chunks: int, dim: int = 0)

def chunk(tensor: Tensor, chunks: int, dim: int = 0) -> 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 chunking it along the dimension 'dim'. It produces 'chunks'
    sub-tensors.

    That operation is only defined for static tensors (no dynamic dimension)
    and the size of the tensor in the dimension 'dim' must be a multiple of
    'chunks': 'input.shape[dim] % chunks == 0'.

    It maps to 'split' with 'split_size = input.shape[dim] / chunks'.

    Parameters:
        tensor : Tensor
            The input tensor to slice.

        chunks : int
            The number of slices to split the input tensor into.

        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]
    assert dim_value % chunks == 0

    return split(tensor, dim_value // chunks, dim)


def unbind(input: Tensor, dim: int = 0):