Versioned name: TopK-1

Category: Sorting and maximization

Short description: TopK computes indices and values of the k maximum/minimum values for each slice along specified axis.


  • axis
    • Description: Specifies the axis along which the values are retrieved.
    • Range of values: An integer. Negative value means counting dimension from the end.
    • Type: int
    • Default value: None
    • Required: yes
  • mode
    • Description: Specifies which operation is used to select the biggest element of two.
    • Range of values: min, max
    • Type: string
    • Default value: None
    • Required: yes
  • sort
    • Description: Specifies order of output elements and/or indices.
    • Range of values: value, index, none
    • Type: string
    • Default value: None
    • Required: yes
  • index_element_type
    • Description: the type of output tensor with indices
    • Range of values: "i64" or "i32"
    • Type: string
    • Default value: "i32"
    • Required: No


  • 1: Arbitrary tensor. Required.
  • 2: k – scalar specifies how many maximum/minimum elements should be computed


  • 1: Output tensor with top k values from the input tensor along specified dimension axis. The shape of the tensor is [input1.shape[0], ..., input1.shape[axis-1], k, input1.shape[axis+1], ...].
  • 2: Output tensor with top k indices for each slice along axis dimension. It is 1D tensor of shape [k]. The shape of the tensor is the same as for the 1st output, that is [input1.shape[0], ..., input1.shape[axis-1], k, input1.shape[axis+1], ...]

Detailed Description

Output tensor is populated by values computes in the following way:

output[i1, ..., i(axis-1), j, i(axis+1) ..., iN] = top_k(input[i1, ...., i(axis-1), :, i(axis+1), ..., iN]), k, sort, mode)

So for each slice input[i1, ...., i(axis-1), :, i(axis+1), ..., iN] which represents 1D array, top_k value is computed individually.

Sorting and minimum/maximum are controlled by sort and mode attributes:

  • mode=max, sort=value - descending by value
  • mode=max, sort=index - ascending by index
  • mode=max, sort=none - undefined
  • mode=min, sort=value - ascending by value
  • mode=min, sort=index - ascending by index
  • mode=min, sort=none - undefined

If there are several elements with the same value then their output order is not determined.


<layer ... type="TopK" ... >
<data axis="1" mode="max" sort="value"/>
<port id="0">
<port id="1">
<!-- k = 3 -->
<port id="2">