TopK

Versioned name: TopK-3

Category: Sorting and maximization

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

Attributes

Inputs:

Outputs:

Types

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, TopK value is computed individually. Sorting and minimum/maximum are controlled by sort and mode attributes.

Example

<layer ... type="TopK" ... >
<data axis="1" mode="max" sort="value" index_element_type="i64"/>
<input>
<port id="0">
<dim>6</dim>
<dim>12</dim>
<dim>10</dim>
<dim>24</dim>
</port>
<port id="1">
</port>
<output>
<port id="2">
<dim>6</dim>
<dim>3</dim>
<dim>10</dim>
<dim>24</dim>
</port>
<port id="3">
<dim>6</dim>
<dim>3</dim>
<dim>10</dim>
<dim>24</dim>
</port>
</output>
</layer>