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

00Retrieve the top-K elements along a specified axis. Given an input tensor ofRetrieve the top-K largest or smallest elements along a specified axis. Given an input tensor of
11shape [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:
22 -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]
33 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
44 -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
55 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
66 tensor). tensor).
77
8If "largest" is 1 (the default value) then the k largest elements are returned.
9If "sorted" is 1 (the default value) then the resulting k elements will be sorted.
10If "sorted" is 0, order of returned 'Values' and 'Indices' are undefined.
11
812Given two equivalent values, this operator uses the indices along the axis asGiven two equivalent values, this operator uses the indices along the axis as
913 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.
1014
1115**Attributes****Attributes**
1216
1317* **axis**:* **axis**:
1418 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.
1525
1626**Inputs****Inputs**
1727
1828* **X** (heterogeneous) - **T**:* **X** (heterogeneous) - **T**:
1929 Tensor of shape [a_1, a_2, ..., a_n, r] Tensor of shape [a_1, a_2, ..., a_n, r]
2030* **K** (heterogeneous) - **tensor(int64)**:* **K** (heterogeneous) - **tensor(int64)**:
2131 A 1-D tensor containing a single positive value corresponding to the A 1-D tensor containing a single positive value corresponding to the
2232 number of top elements to retrieve number of top elements to retrieve
2333
2434**Outputs****Outputs**
2535
2636* **Values** (heterogeneous) - **T**:* **Values** (heterogeneous) - **T**:
2737 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]
2838 containing top K values from the input tensor containing top K values from the input tensor
2939* **Indices** (heterogeneous) - **I**:* **Indices** (heterogeneous) - **I**:
3040 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]
3141 containing the corresponding input tensor indices for the top K containing the corresponding input tensor indices for the top K
3242 values. values.
3343
3444**Type Constraints****Type Constraints**
3545
3646* **T** in (* **T** in (
3747 tensor(double), tensor(double),
3848 tensor(float), tensor(float),
3949 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)
4058 ): ):
4159 Constrain input and output types to float tensors. Constrain input and output types to numeric tensors.
4260* **I** in (* **I** in (
4361 tensor(int64) tensor(int64)
4462 ): ):
4563 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

00Retrieve the top-K elements along a specified axis. Given an input tensor ofRetrieve the top-K elements along a specified axis. Given an input tensor of
11shape [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:
22 -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]
33 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
44 -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
55 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
66 tensor). tensor).
7
78Given two equivalent values, this operator uses the indices along the axis asGiven two equivalent values, this operator uses the indices along the axis as
89 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.
910
1011**Attributes****Attributes**
1112
1213* **axis**:* **axis**:
1314 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
1615
1716**Inputs****Inputs**
1817
1918* **X** (heterogeneous) - **T**:* **X** (heterogeneous) - **T**:
2019 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
2123
2224**Outputs****Outputs**
2325
2426* **Values** (heterogeneous) - **T**:* **Values** (heterogeneous) - **T**:
2527 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]
2628 containing top K values from the input tensor containing top K values from the input tensor
2729* **Indices** (heterogeneous) - **I**:* **Indices** (heterogeneous) - **I**:
2830 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]
2931 containing the corresponding input tensor indices for the top K containing the corresponding input tensor indices for the top K
3032 values. values.
3133
3234**Type Constraints****Type Constraints**
3335
3436* **T** in (* **T** in (
3537 tensor(double), tensor(double),
3638 tensor(float), tensor(float),
3739 tensor(float16) tensor(float16)
3840 ): ):
3941 Constrain input and output types to float tensors. Constrain input and output types to float tensors.
4042* **I** in (* **I** in (
4143 tensor(int64) tensor(int64)
4244 ): ):
4345 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