Function expand_dims_like

tensorrt_llm/functional.py:2010–2051

Add an operation to expand the first tensor to the same rank as the second tensor. That function takes a first tensor. It also accepts an integer or a float, in which case it creates a constant tensor from it. In both cases, the rank of that first tensor is compared to the rank

(left: Union[Tensor, int, float], right: Tensor)

def expand_dims_like(left: Union[Tensor, int, float], right: Tensor) -> Tensor:
    '''
    Add an operation to expand the first tensor to the same rank as the second
    tensor.

    That function takes a first tensor. It also accepts an integer or a float,
    in which case it creates a constant tensor from it. In both cases, the rank
    of that first tensor is compared to the rank of the second tensor. If they
    are of the same rank, the first tensor is returned. Otherwise, the first
    tensor is expanded on the left to match the rank of the second tensor.

    Note that the shapes do not have to match, only the rank is considered in
    that function.

    For example, for a pair of tensors of shapes [3, 4] and [4, 3, 2], the
    first tensor will be expanded to a tensor of rank 3 and shape [1, 3, 4].

    Parameters:
        left : Union[Tensor, int, float]
            The first tensor to expand. When a scalar value is provided as a
            parameter, that function first creates a tensor before expanding it
            (if needed).

        right : Tensor
            The reference tensor to match.

    Returns:
        The tensor produced by the shuffle layer.
    '''
    if isinstance(left, int):
        left = constant(dims_array([left]))
    elif isinstance(left, float):
        if isinstance(right, Tensor) and right.dtype == trt.DataType.HALF:
            left = constant(fp16_array([left]))
        else:
            left = constant(fp32_array([left]))
    left_ndim = left.ndim()
    right_ndim = right.ndim()
    if right_ndim > left_ndim:
        new_ndim = list(range(right_ndim - left_ndim))
        return expand_dims(left, new_ndim)
    return left


# If dim is None, return a 1-D TensorRT LLM tensor of the size