Function stack

tensorrt_llm/functional.py:1961–2007

Add an operation to contact input tensors along a new dimension. The function creates an operation that creates a new dim for all the input tensors and then concatenates them along that new dim. . All the tensors in 'inputs' must have the same shape. for ii in range(input

(inputs: Sequence[Tensor], dim: int = 0)

def stack(inputs: Sequence[Tensor], dim: int = 0) -> Tensor:
    '''
    Add an operation to contact input tensors along a new dimension.

    The function creates an operation that creates a new dim for all the
    input tensors and then concatenates them along that new dim.
.

    All the tensors in 'inputs' must have the same shape.

        for ii in range(inputs[0].rank()):
            assert all(inp.shape[ii] == inputs[0].shape[ii] for inp in inputs)

    The shape of the output tensor is defined as:

        output.rank() = inputs[0].rank() + 1

        output.shape[dim] = len(inputs)

        for ii in range(inputs[0].rank()):
            if ii < dim:
                output.shape[ii] = inputs[0].shape[ii]
            else:
                output.shape[ii+1] = inputs[0].shape[ii]

    For example, given a sequence of two 2D tensors [[0, 1], [2, 3]] and
    [[4, 5], [6, 7]] both of shape [2, 2],

        stack(inputs, 0)

    will produce [[[0, 1], [2, 3]], [[4, 5], [6, 7]]] of shape [2, 2, 2] and

        stack(inputs, 1)

    will produce [[[0, 1], [4, 5]], [[2, 3], [6, 7]]] of shape [2, 2, 2].

    Parameters:
        inputs : Sequence[Tensor]
            The sequence of tensors to stack.

        dim : int
            The dimension in which the stack is performed.

    Returns:
        A tensor that contains the input tensors stacked along a new dimension.
    ''&#x27;
    return concat([unsqueeze(inp, axis=dim) for inp in inputs], dim=dim)


def expand_dims_like(left: Union[Tensor, int, float], right: Tensor) -> Tensor: