TopK

Description

Obtains the first K largest or smallest elements along the specified axis of the input tensor.

Input

X: tensor of type float16 or float.

K: 1D tensor of type int64, specifying the first K elements to be obtained.

Attribute

axis: int, axis for sorting.

largest: int, specifying the first K largest or smallest values to be returned. If the value is 1, the first K largest values are returned. If the value is 0, the first K smallest values are returned.

sorted: int, specifying whether to sort the output. If the value is 1, the output is sorted in ascending order. If the value is 0, the output is not sorted.

Output

Values: tensor, which is the first K largest elements on a given axis. The data type is the same as X, and shape is [a_0, a_1, ..., a_axis-1, K, a_axis+1, ...].

Indices: tensor of type int64, which is the indices of the first K largest elements. The shape is the same as that of Values.

Constraints

None.

ONNX Opset Support

Opset v8/v9/v10/v11/v12/v13/v14/v15/v16/v17/v18