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).
largest: Whether to return the top-K largest or smallest elements.
sorted: Whether to return the elements in sorted order.
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
import numpy as np
import onnx
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
import numpy as np
import onnx
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
import numpy as np
import onnx
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",
)
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.
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
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.
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