Function masked_select

tensorrt_llm/functional.py:2352–2408

Add an operation to select elements from a tensor according to a boolean mask tensor. Given an input tensor, that function creates an operation that selects elements at the indices indicated by the boolean mask tensor to create a new tensor. The output tensor is a 1-D tensor.

(input: Tensor, mask: Tensor)

def masked_select(input: Tensor, mask: Tensor) -> Tensor:
    '''
    Add an operation to select elements from a tensor according to a boolean
    mask tensor.

    Given an input tensor, that function creates an operation that selects
    elements at the indices indicated by the boolean mask tensor to create
    a new tensor. The output tensor is a 1-D tensor.

    The input tensor must have rank >= 1. The shapes of the input tensor and
    the mask tensor don’t need to match, but they must be able to be broadcasted.

    For example, on input=[[4, 2, 5], [2, 1, 2], [4, 7, 1]], which has a shape
    [3, 3],

        masked_select(input, [[True, False, True], [False, True, False], [True, False, True]])

    will create a tensor of shape [5] that contains the [4, 5, 1, 4, 1].

        masked_select(input, [[True], [False], [True]])

    will create a tensor of shape [6] that contains the [4, 2, 5, 4, 7, 1].

        masked_select(input, [[False, False, True]])

    will create a tensor of shape [3] that contains the [5, 2, 1].

        masked_select(input, [False])

    will create a tensor of shape [0] which is empty.

    That operation is implemented by NonZero, Shuffle and GatherV2 layers
    in TensorRT.

    Parameters:
        input : Tensor
            The input tensor to select from.

        mask : Tensor
            The boolean mask tensor that indicates elements to select.

    Returns:
        The 1-D tensor containing the selected elements.
    '''
    assert input.rank() >= 1, "input should have rank >= 1"
    input, mask = broadcast_helper(input, mask)
    expanded_mask = expand(mask, shape(input))

    non_zero_layer = default_trtnet().add_non_zero(expanded_mask.trt_tensor)

    shuffle_layer = default_trtnet().add_shuffle(non_zero_layer.get_output(0))
    shuffle_layer.second_transpose = (1, 0)

    gather_layer = default_trtnet().add_gather_v2(input.trt_tensor,
                                                  shuffle_layer.get_output(0),
                                                  mode=trt.GatherMode.ND)
    return _create_tensor(gather_layer.get_output(0), gather_layer)