Function where

tensorrt_llm/functional.py:3043–3099

Add a where (aka select or if-then-else) operation. Assuming the three input parameters have the same shape, that function creates the operation to compute a tensor of the same shape such that: for ii in range(mul(condition.shape)): output[ii] = left[ii] if conditi

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

def where(condition: Union[Tensor, bool], left: Union[Tensor, int, float],
          right: Union[Tensor, int, float]) -> Tensor:
    '''
    Add a where (aka select or if-then-else) operation.

    Assuming the three input parameters have the same shape, that function creates
    the operation to compute a tensor of the same shape such that:

        for ii in range(mul(condition.shape)):
            output[ii] = left[ii] if condition[ii] else right[ii]

    For each input, that function first creates a constant tensor if the
    condition is boolean or the left/right 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. Then, it performs the selection.

    It is implemented using the ISelectLayer from TensorRT.

    Parameters:
        condition : Union[Tensor, bool]
            The condition. If that input is a boolean, the function
            creates a constant tensor.

        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:
        The tensor produced by this where operation.
    '''
    # Convert to tensors.
    condition = constant_to_tensor_(condition)
    left, right = constants_to_tensors_(left, right)

    # Find the tensor with the largest rank of the three.
    largest = condition
    if largest.rank() < left.rank():
        largest = left
    if largest.rank() < right.rank():
        largest = right

    # Expand the tensors to match the largest one.
    if condition is not largest:
        condition = expand_dims_like(condition, largest)
    if left is not largest:
        left = expand_dims_like(left, largest)
    if right is not largest:
        right = expand_dims_like(right, largest)

    # Insert the operation.
    layer = default_trtnet().add_select(condition.trt_tensor, left.trt_tensor,
                                        right.trt_tensor)
    return _create_tensor(layer.get_output(0), layer)