Function expand

tensorrt_llm/functional.py:1572–1627

Add an operation to expand a tensor. The operation expands the input tensor in the singleton dimensions to the size indicated by the corresponding dimension in the `expand_shape` tensor. In other words, given an input tensor with dimensions of size 1, those dimensions will be e

(input: Tensor, expand_shape: Tensor)

def expand(input: Tensor, expand_shape: Tensor) -> Tensor:
    '''
    Add an operation to expand a tensor.

    The operation expands the input tensor in the singleton dimensions to the
    size indicated by the corresponding dimension in the `expand_shape` tensor.
    In other words, given an input tensor with dimensions of size 1, those
    dimensions will be expanded to the size in `expand_shape`.

    For example, a tensor of shape [4, 3, 1, 3] will be expanded to a tensor of
    shape [4, 3, 2, 3] by the layer created using expand(input, [4, 3, 2, 3]).

    The expansion may either replicate the values or be mapped to a view with a
    stride of 0 in the expanded dimensions. For example, for a tensor [[3, 2]] of
    shape [1, 2],

        expand([[3, 2]], [2, 2])

    can be used to expand the input to [[3, 2], [3, 2]].

    This operation is implemented using a tensorrt.ISliceLayer. The current
    implementation does not verify that non singleton dimensions are not
    shrunk. In other words, for an input of shape [4, 1, 2],

        expand(input, [3, 2, 2])

    will produce a tensor of shape [3, 2, 2]. That behavior is subject to
    change in the future.

    Parameters:
        input : Tensor
            The input tensor.

        expand_shape : Tensor
            The new shape of the expanded tensor.

    Returns:
        The tensor produced by the expand layer.
    '''
    ndim = input.rank()
    layer = default_trtnet().add_slice(
        input.trt_tensor,
        start=[0 for _ in range(ndim)],
        shape=[1 for _ in range(ndim)],  # unused dummy value
        stride=[1 for _ in range(ndim)]  # unused dummy value
    )

    # The stride is either:
    #   0 for dimensions of size 1 (i.e. shape(input, i) - 1 == 1 - 1 == 0) or,
    #   1 for dimensions of size > 1 since minimum(value >= 1, 1) == 1.
    stride_tensor = concat(
        [minimum((shape(input, i) - 1), 1) for i in range(ndim)])

    layer.set_input(2, expand_shape.trt_tensor)
    layer.set_input(3, stride_tensor.trt_tensor)
    return _create_tensor(layer.get_output(0), layer)