TopK¶

TopK - 24¶

Version¶

name: TopK (GitHub)
domain: main
since_version: 24
function: False
support_level: SupportType.COMMON
shape inference: True

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

Summary¶

Retrieve the top-K largest or smallest elements along a specified axis. Given an input tensor of shape [a_0, a_1, …, a_{n-1}] and integer argument k, return two outputs:

Value tensor of shape [a_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] which contains the values of the top k elements along the specified axis
Index tensor of shape [a_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] 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 - INT (default is -1):

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 - INT (default is 1):

Whether to return the top-K largest or smallest elements.
sorted - INT (default is 1):

Whether to return the elements in sorted order.

Inputs¶

X (heterogeneous) - T:

Tensor of shape [a_0, a_1, …, a_{n-1}]
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_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] containing top K values from the input tensor
Indices (heterogeneous) - I:

Tensor of shape [a_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] containing the corresponding input tensor indices for the top K values.

Type Constraints¶

T in ( tensor(bfloat16), 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_uint64¶

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.uint64,
)
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_uint64",
)

_top_k_same_values¶

import numpy as np
import onnx

axis = 0
largest = 0

k = 3
node = onnx.helper.make_node(
    "TopK", inputs=["x", "k"], outputs=["values", "indices"], axis=axis
)
X = np.array(
    [0, 0, 0, 0],
    dtype=np.int64,
)
K = np.array([k], dtype=np.int64)
values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)

# (Pdb) print(values_ref)
# [0 0 0]
# (Pdb) print(indices_ref)
# [0 1 2]

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

_top_k_same_values_largest¶

import numpy as np
import onnx

axis = 0
largest = 1

k = 3
node = onnx.helper.make_node(
    "TopK", inputs=["x", "k"], outputs=["values", "indices"], axis=axis
)
X = np.array(
    [0, 0, 0, 0],
    dtype=np.int64,
)
K = np.array([k], dtype=np.int64)
values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)

# print(values_ref)
# [0 0 0]
# print(indices_ref)
# [0 1 2]

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

_top_k_same_values_2d¶

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, 0, 0, 0], [1, 1, 1, 1], [2, 2, 1, 1]],
    dtype=np.int64,
)
K = np.array([k], dtype=np.int64)
values_ref, indices_ref = topk_sorted_implementation(X, k, axis, largest)

# print(values_ref)
# [[0 0 0]
# [1 1 1]
# [1 1 2]]
# print(indices_ref)
# [[0 1 2]
# [0 1 2]
# [2 3 0]]

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

_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 - 11 vs 24

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_0, a_1, …, a_{n-1}] and integer argument k, return two outputs:

Value tensor of shape [a_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] which contains the values of the top k elements along the specified axis
Index tensor of shape [a_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] 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 - INT (default is -1):

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 - INT (default is 1):

Whether to return the top-K largest or smallest elements.
sorted - INT (default is 1):

Whether to return the elements in sorted order.

Inputs¶

X (heterogeneous) - T:

Tensor of shape [a_0, a_1, …, a_{n-1}]
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_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] containing top K values from the input tensor
Indices (heterogeneous) - I:

Tensor of shape [a_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] 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

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_0, a_1, …, a_{n-1}] and integer argument k, return two outputs: -Value tensor of shape [a_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] which contains the values of the top k elements along the specified axis -Index tensor of shape [a_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] 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 - INT (default is -1):

Dimension on which to do the sort.

Inputs¶

X (heterogeneous) - T:

Tensor of shape [a_0, a_1, …, a_{n-1}]
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_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] containing top K values from the input tensor
Indices (heterogeneous) - I:

Tensor of shape [a_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] 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¶

Attributes¶

axis - INT (default is -1):

Dimension on which to do the sort.
k - INT (required) :

Number of top elements to retrieve

Inputs¶

X (heterogeneous) - T:

Tensor of shape [a_0, a_1, …, a_{n-1}]

Outputs¶

Values (heterogeneous) - T:

Tensor of shape [a_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] containing top K values from the input tensor
Indices (heterogeneous) - I:

Tensor of shape [a_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] 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