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