Function gather_nd

tensorrt_llm/functional.py:2313–2333

Adds a layer that performs a gather with some element-wise dimensions. See: https://onnx.ai/onnx/operators/onnx__GatherND.html The gather is performed on dim=batch_dims. Parameters: input: Tensor The tensor on which the gather operation is performed. ind

(input: Tensor, indices: Tensor, batch_dims: int = 1)

def gather_nd(input: Tensor, indices: Tensor, batch_dims: int = 1) -> Tensor:
    '''
    Adds a layer that performs a gather with some element-wise dimensions.
    See: https://onnx.ai/onnx/operators/onnx__GatherND.html
    The gather is performed on dim=batch_dims.

    Parameters:
        input: Tensor
            The tensor on which the gather operation is performed.
        indices: Tensor
            The tensor that indicates which entries to be gathered.
        batch_dims: int
            The number of first dimensions that should be skipped before gather starts.
    Returns:
        A tensor created by the gather layer with GatherMode.ND.
    '''
    gather_layer = default_trtnet().add_gather_v2(input.trt_tensor,
                                                  indices.trt_tensor,
                                                  mode=trt.GatherMode.ND)
    gather_layer.num_elementwise_dims = batch_dims
    return _create_tensor(gather_layer.get_output(0), gather_layer)


def nonzero(input: Tensor) -> Tensor: