TopK#

TopK - 11#

Version

• name: TopK (GitHub)

• domain: main

• since_version: 11

• function: False

• support_level: SupportType.COMMON

• shape inference: True

This version of the operator has been available since version 11.

Summary

Retrieve the top-K largest or smallest elements along a specified axis. Given an input tensor of shape [a_1, a_2, …, a_n, r] and integer argument k, return two outputs:

-Value tensor of shape [a_1, a_2, …, a_{axis-1}, k, a_{axis+1}, … a_n]

which contains the values of the top k elements along the specified axis

-Index tensor of shape [a_1, a_2, …, a_{axis-1}, k, a_{axis+1}, … a_n] which

contains the indices of the top k elements (original indices from the input tensor).

If “largest” is 1 (the default value) then the k largest elements are returned. If “sorted” is 1 (the default value) then the resulting k elements will be sorted. If “sorted” is 0, order of returned ‘Values’ and ‘Indices’ are undefined.

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.

Attributes

• axis: Dimension on which to do the sort. Negative value means counting dimensions from the back. Accepted range is [-r, r-1] where r = rank(input). Default value is `-1`.

• largest: Whether to return the top-K largest or smallest elements. Default value is `1`.

• sorted: Whether to return the elements in sorted order. Default value is `1`.

Inputs

• X (heterogeneous) - T: Tensor of shape [a_1, a_2, …, a_n, r]

• K (heterogeneous) - tensor(int64): A 1-D tensor containing a single positive value corresponding to the number of top elements to retrieve

Outputs

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

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

Type Constraints

• T in ( tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8) ): Constrain input and output types to numeric tensors.

• I in ( tensor(int64) ): Constrain index tensor to int64

Examples

_top_k

```axis = 1
largest = 1

k = 3
node = onnx.helper.make_node(
"TopK", inputs=["x", "k"], outputs=["values", "indices"], axis=axis
)
X = np.array(
[
[0, 1, 2, 3],
[4, 5, 6, 7],
[8, 9, 10, 11],
],
dtype=np.float32,
)
K = np.array([k], dtype=np.int64)
values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)

# print(values_ref)
# [[ 3.  2.  1.]
# [ 7.  6.  5.]
# [11. 10.  9.]]
# print(indices_ref)
# [[3 2 1]
# [3 2 1]
# [3 2 1]]

expect(
node, inputs=[X, K], outputs=[values_ref, indices_ref], name="test_top_k"
)
```

_top_k_smallest

```axis = 1
largest = 0
sorted = 1
k = 3

node = onnx.helper.make_node(
"TopK",
inputs=["x", "k"],
outputs=["values", "indices"],
axis=axis,
largest=largest,
sorted=sorted,
)

X = np.array(
[
[0, 1, 2, 3],
[4, 5, 6, 7],
[11, 10, 9, 8],
],
dtype=np.float32,
)
K = np.array([k], dtype=np.int64)
values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)

# print(values_ref)
# [[ 0.  1.  2.]
# [ 4.  5.  6.]
# [ 8.  9. 10.]]
# print(indices_ref)
# [[0 1 2]
# [0 1 2]
# [3 2 1]]

expect(
node,
inputs=[X, K],
outputs=[values_ref, indices_ref],
name="test_top_k_smallest",
)
```

_top_k_negative_axis

```axis = -1
largest = 1

k = 3
node = onnx.helper.make_node(
"TopK", inputs=["x", "k"], outputs=["values", "indices"], axis=axis
)
X = np.array(
[
[0, 1, 2, 3],
[4, 5, 6, 7],
[8, 9, 10, 11],
],
dtype=np.float32,
)
K = np.array([k], dtype=np.int64)
values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)

# print(values_ref)
# [[ 3.  2.  1.]
# [ 7.  6.  5.]
# [11. 10.  9.]]
# print(indices_ref)
# [[3 2 1]
# [3 2 1]
# [3 2 1]]

expect(
node,
inputs=[X, K],
outputs=[values_ref, indices_ref],
name="test_top_k_negative_axis",
)
```

Differences

 `0` `0` `Retrieve the top-K elements along a specified axis. Given an input tensor of` `Retrieve the top-K largest or smallest elements along a specified axis. Given an input tensor of` `1` `1` `shape [a_1, a_2, ..., a_n, r] and integer argument k, return two outputs:` `shape [a_1, a_2, ..., a_n, r] and integer argument k, return two outputs:` `2` `2` ` -Value tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]` ` -Value tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]` `3` `3` ` which contains the values of the top k elements along the specified axis` ` which contains the values of the top k elements along the specified axis` `4` `4` ` -Index tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] which` ` -Index tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] which` `5` `5` ` contains the indices of the top k elements (original indices from the input` ` contains the indices of the top k elements (original indices from the input` `6` `6` ` tensor).` ` tensor).` `7` `7` `8` `If "largest" is 1 (the default value) then the k largest elements are returned.` `9` `If "sorted" is 1 (the default value) then the resulting k elements will be sorted.` `10` `If "sorted" is 0, order of returned 'Values' and 'Indices' are undefined.` `11` `8` `12` `Given two equivalent values, this operator uses the indices along the axis as` `Given two equivalent values, this operator uses the indices along the axis as` `9` `13` ` a tiebreaker. That is, the element with the lower index will appear first.` ` a tiebreaker. That is, the element with the lower index will appear first.` `10` `14` `11` `15` `**Attributes**` `**Attributes**` `12` `16` `13` `17` `* **axis**:` `* **axis**:` `14` `18` ` Dimension on which to do the sort. Default value is -1.` ` Dimension on which to do the sort. Negative value means counting` `19` ` dimensions from the back. Accepted range is [-r, r-1] where r =` `20` ` rank(input). Default value is -1.` `21` `* **largest**:` `22` ` Whether to return the top-K largest or smallest elements. Default value is 1.` `23` `* **sorted**:` `24` ` Whether to return the elements in sorted order. Default value is 1.` `15` `25` `16` `26` `**Inputs**` `**Inputs**` `17` `27` `18` `28` `* **X** (heterogeneous) - **T**:` `* **X** (heterogeneous) - **T**:` `19` `29` ` Tensor of shape [a_1, a_2, ..., a_n, r]` ` Tensor of shape [a_1, a_2, ..., a_n, r]` `20` `30` `* **K** (heterogeneous) - **tensor(int64)**:` `* **K** (heterogeneous) - **tensor(int64)**:` `21` `31` ` A 1-D tensor containing a single positive value corresponding to the` ` A 1-D tensor containing a single positive value corresponding to the` `22` `32` ` number of top elements to retrieve` ` number of top elements to retrieve` `23` `33` `24` `34` `**Outputs**` `**Outputs**` `25` `35` `26` `36` `* **Values** (heterogeneous) - **T**:` `* **Values** (heterogeneous) - **T**:` `27` `37` ` Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]` ` Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]` `28` `38` ` containing top K values from the input tensor` ` containing top K values from the input tensor` `29` `39` `* **Indices** (heterogeneous) - **I**:` `* **Indices** (heterogeneous) - **I**:` `30` `40` ` Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]` ` Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]` `31` `41` ` containing the corresponding input tensor indices for the top K` ` containing the corresponding input tensor indices for the top K` `32` `42` ` values.` ` values.` `33` `43` `34` `44` `**Type Constraints**` `**Type Constraints**` `35` `45` `36` `46` `* **T** in (` `* **T** in (` `37` `47` ` tensor(double),` ` tensor(double),` `38` `48` ` tensor(float),` ` tensor(float),` `39` `49` ` tensor(float16)` ` tensor(float16),` `50` ` tensor(int16),` `51` ` tensor(int32),` `52` ` tensor(int64),` `53` ` tensor(int8),` `54` ` tensor(uint16),` `55` ` tensor(uint32),` `56` ` tensor(uint64),` `57` ` tensor(uint8)` `40` `58` ` ):` ` ):` `41` `59` ` Constrain input and output types to float tensors.` ` Constrain input and output types to numeric tensors.` `42` `60` `* **I** in (` `* **I** in (` `43` `61` ` tensor(int64)` ` tensor(int64)` `44` `62` ` ):` ` ):` `45` `63` ` Constrain index tensor to int64` ` Constrain index tensor to int64`

TopK - 10#

Version

• name: TopK (GitHub)

• domain: main

• since_version: 10

• function: False

• support_level: SupportType.COMMON

• shape inference: True

This version of the operator has been available since version 10.

Summary

Retrieve the top-K elements along a specified axis. Given an input tensor of shape [a_1, a_2, …, a_n, r] and integer argument k, return two outputs:

-Value tensor of shape [a_1, a_2, …, a_{axis-1}, k, a_{axis+1}, … a_n]

which contains the values of the top k elements along the specified axis

-Index tensor of shape [a_1, a_2, …, a_{axis-1}, k, a_{axis+1}, … a_n] 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.

Attributes

• axis: Dimension on which to do the sort. Default value is `-1`.

Inputs

• X (heterogeneous) - T: Tensor of shape [a_1, a_2, …, a_n, r]

• K (heterogeneous) - tensor(int64): A 1-D tensor containing a single positive value corresponding to the number of top elements to retrieve

Outputs

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

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

Type Constraints

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

• I in ( tensor(int64) ): Constrain index tensor to int64

Differences

 `0` `0` `Retrieve the top-K elements along a specified axis. Given an input tensor of` `Retrieve the top-K elements along a specified axis. Given an input tensor of` `1` `1` `shape [a_1, a_2, ..., a_n, r] and integer argument k, return two outputs:` `shape [a_1, a_2, ..., a_n, r] and integer argument k, return two outputs:` `2` `2` ` -Value tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]` ` -Value tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]` `3` `3` ` which contains the values of the top k elements along the specified axis` ` which contains the values of the top k elements along the specified axis` `4` `4` ` -Index tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] which` ` -Index tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] which` `5` `5` ` contains the indices of the top k elements (original indices from the input` ` contains the indices of the top k elements (original indices from the input` `6` `6` ` tensor).` ` tensor).` `7` `7` `8` `Given two equivalent values, this operator uses the indices along the axis as` `Given two equivalent values, this operator uses the indices along the axis as` `8` `9` ` a tiebreaker. That is, the element with the lower index will appear first.` ` a tiebreaker. That is, the element with the lower index will appear first.` `9` `10` `10` `11` `**Attributes**` `**Attributes**` `11` `12` `12` `13` `* **axis**:` `* **axis**:` `13` `14` ` Dimension on which to do the sort. Default value is -1.` ` Dimension on which to do the sort. Default value is -1.` `14` `* **k** (required):` `15` ` Number of top elements to retrieve` `16` `15` `17` `16` `**Inputs**` `**Inputs**` `18` `17` `19` `18` `* **X** (heterogeneous) - **T**:` `* **X** (heterogeneous) - **T**:` `20` `19` ` Tensor of shape [a_1, a_2, ..., a_n, r]` ` Tensor of shape [a_1, a_2, ..., a_n, r]` `20` `* **K** (heterogeneous) - **tensor(int64)**:` `21` ` A 1-D tensor containing a single positive value corresponding to the` `22` ` number of top elements to retrieve` `21` `23` `22` `24` `**Outputs**` `**Outputs**` `23` `25` `24` `26` `* **Values** (heterogeneous) - **T**:` `* **Values** (heterogeneous) - **T**:` `25` `27` ` Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]` ` Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]` `26` `28` ` containing top K values from the input tensor` ` containing top K values from the input tensor` `27` `29` `* **Indices** (heterogeneous) - **I**:` `* **Indices** (heterogeneous) - **I**:` `28` `30` ` Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]` ` Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]` `29` `31` ` containing the corresponding input tensor indices for the top K` ` containing the corresponding input tensor indices for the top K` `30` `32` ` values.` ` values.` `31` `33` `32` `34` `**Type Constraints**` `**Type Constraints**` `33` `35` `34` `36` `* **T** in (` `* **T** in (` `35` `37` ` tensor(double),` ` tensor(double),` `36` `38` ` tensor(float),` ` tensor(float),` `37` `39` ` tensor(float16)` ` tensor(float16)` `38` `40` ` ):` ` ):` `39` `41` ` Constrain input and output types to float tensors.` ` Constrain input and output types to float tensors.` `40` `42` `* **I** in (` `* **I** in (` `41` `43` ` tensor(int64)` ` tensor(int64)` `42` `44` ` ):` ` ):` `43` `45` ` Constrain index tensor to int64` ` Constrain index tensor to int64`

TopK - 1#

Version

• name: TopK (GitHub)

• domain: main

• since_version: 1

• function: False

• support_level: SupportType.COMMON

• shape inference: True

This version of the operator has been available since version 1.

Summary

Retrieve the top-K elements along a specified axis. Given an input tensor of shape [a_1, a_2, …, a_n, r] and integer argument k, return two outputs:

-Value tensor of shape [a_1, a_2, …, a_{axis-1}, k, a_{axis+1}, … a_n]

which contains the values of the top k elements along the specified axis

-Index tensor of shape [a_1, a_2, …, a_{axis-1}, k, a_{axis+1}, … a_n] 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.

Attributes

• axis: Dimension on which to do the sort. Default value is `-1`.

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

Inputs

• X (heterogeneous) - T: Tensor of shape [a_1, a_2, …, a_n, r]

Outputs

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

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

Type Constraints

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

• I in ( tensor(int64) ): Constrain index tensor to int64