Function nonzero

tensorrt_llm/functional.py:2336–2349

Adds a layer that finds the indices of non-zero values of the input tensor. Parameters: input: Tensor The input tensor for which we need to find the indices of non-zero values. Returns: A tensor of shape [D, C] where D is the number of dimensions of `input`

(input: Tensor)

def nonzero(input: Tensor) -> Tensor:
    '''
    Adds a layer that finds the indices of non-zero values of the input tensor.

    Parameters:
        input: Tensor
            The input tensor for which we need to find the indices of non-zero values.
    Returns:
        A tensor of shape [D, C] where D is the number of dimensions of `input` and
        C is the number of non-zero values in it.
        Each column of this 2D tensor represents the index tuple for each non-zero value.
    '''
    non_zero_layer = default_trtnet().add_non_zero(input.trt_tensor)
    return _create_tensor(non_zero_layer.get_output(0), non_zero_layer)


def masked_select(input: Tensor, mask: Tensor) -> Tensor: