Function broadcast_helper

tensorrt_llm/functional.py:2920–2958

Helper function to perform a broadcast. For each input, that function first creates a constant tensor if the input is an integer or a float. Then, if needed, it expands the smaller tensor to make sure its rank is the same as the larger one. Parameters: left : Union[Ten

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

def broadcast_helper(left: Union[Tensor, int, float],
                     right: Union[Tensor, int, float]) -> Tuple[Tensor, Tensor]:
    '''
    Helper function to perform a broadcast.

    For each input, that function first creates a constant tensor if the input
    is an integer or a float. Then, if needed, it expands the smaller tensor to
    make sure its rank is the same as the larger one.

    Parameters:
        left : Union[Tensor, int, float]
            The first input. If that input is an integer or a float, the
            function creates a constant tensor.

        right : Union[Tensor, int, float]
            The second input. If that input is an integer or a float, the
            function creates a constant tensor.

    Returns:
        A pair of tensors of same rank.
    '''
    if not default_net().strongly_typed:
        left = constant_to_tensor_(left)
        right = constant_to_tensor_(right)
    else:
        left = constant_to_tensor_(
            left, right.dtype if isinstance(right, Tensor) else None)
        right = constant_to_tensor_(right, left.dtype)

    if left.rank() == right.rank():
        return (left, right)

    if left.rank() < right.rank():
        left = expand_dims_like(left, right)
        return (left, right)

    if left.rank() > right.rank():
        right = expand_dims_like(right, left)
        return (left, right)


def elementwise_binary(left: Union[Tensor, int,