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.
Attributes
- 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
Inputs:
- 1: Arbitrary tensor. Required.
- 2: k – scalar specifies how many maximum/minimum elements should be computed
Outputs:
- 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.
Example
<layer ... type="TopK" ... >
<data axis="1" mode="max" sort="value"/>
<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>
</output>
</layer>