TopK

Retrieves the top k elements from an input tensor and outputs a tensor with those values and another tensor with their indices.

Attributes

operation TopK operation can be one of:

• MIN Retrieves the top minimum elements.

• MAX Retrieves the top maximum elements.

axes the axis to sample; must be one of the last four dimensions.

k The number of elements to retrieve. For axis with dimension size \(d\), k must obey \(k<=d, k<=3840\).

Inputs

input: tensor of type T1.

Outputs

values: tensor of type T1.

indices: tensor of type T2.

Data Types

T1: int32, int64, float16, float32, bfloat16

T2: int32

Shape Information

input is a tensor with a shape of \([a_0,...,a_n]\)

values is a tensor of shape \([a_0,...,a_{axes-1},k,a_{axes+1},...,a_n]\)

indices is a tensor of shape \([a_0,...,a_{axes-1},k,a_{axes+1},...,a_n]\)

Examples

TopK

in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 5))
layer = network.add_topk(in1, op=trt.TopKOperation.MAX, k=3, axes=2)
network.mark_output(layer.get_output(0))
network.mark_output(layer.get_output(1))

inputs[in1.name] = np.array([[-3.0, -2.0, -1.0, 10.0, -25.0], [0.0, 1.0, 2.0, -2.0, -1.0]])

outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[10.0, -1.0, -2.0], [2.0, 1.0, 0.0]])

outputs[layer.get_output(1).name] = layer.get_output(1).shape
expected[layer.get_output(1).name] = np.array([[3, 2, 1], [2, 1, 0]])

C++ API

For more information about the C++ ITopKLayer operator, refer to the C++ ITopKLayer documentation.

Python API

For more information about the Python ITopKLayer operator, refer to the Python ITopKLayer documentation.