TopK - version 1#

This page documents version 1 of operator TopK. See TopK for the latest version (since version 11).

  • Domain: ai.onnx

  • Since version: 1

Retrieve the top-K elements along a specified axis. Given an input tensor of shape [a_0, a_1, …, a{n-1}] and integer argument k, return two outputs:

-Value tensor of shape [a_0, a_1, ..., a_{axis-1}, k, a_{axis+1}, ... a_{n-1}]
  which contains the values of the top k elements along the specified axis
-Index tensor of shape [a_0, a_1, ..., a_{axis-1}, k, a_{axis+1}, ... a_{n-1}] which
 contains the indices of the top k elements (original indices from the input
 tensor).

Given two equivalent values, this operator uses the indices along the axis as

a tiebreaker. That is, the element with the lower index will appear first.

Inputs

  • X (T): Tensor of shape [a_0, a_1, …, a{n-1}]

Outputs

  • Values (T): Tensor of shape [a_0, a_1, …, a{axis-1}, k, a{axis+1}, … a{n-1}] containing top K values from the input tensor

  • Indices (I): Tensor of shape [a_0, a_1, …, a{axis-1}, k, a{axis+1}, … a{n-1}] containing the corresponding input tensor indices for the top K values.

Attributes

  • axis (int): Dimension on which to do the sort.

  • k (int): Number of top elements to retrieve

Type Constraints

  • T: Constrain input and output types to float tensors. Allowed types: tensor(double), tensor(float), tensor(float16).

  • I: Constrain index tensor to int64 Allowed types: tensor(int64).