TopK#

  • Domain: ai.onnx

  • Since version: 11

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) then the k largest elements are returned.

  • If “sorted” is 1 (the default) 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.

Inputs

  • X (T): Tensor of shape [a_0, a_1, …, a{n-1}]

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

Outputs

  • Values (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 (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.

Attributes

  • axis (int): 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): Whether to return the top-K largest or smallest elements.

  • sorted (int): Whether to return the elements in sorted order.

Type Constraints

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

  • I: Constrain index tensor to int64 Allowed types: tensor(int64).

Examples#

test_cc_top_k

Node:
  TopK(x, k) -> (values, indices)
  Attributes:
    axis = 1

Inputs:
  x: shape=(3, 4), dtype=float32
    [[ 0.,  1.,  2.,  3.],
     [ 4.,  5.,  6.,  7.],
     [ 8.,  9., 10., 11.]]
  k: shape=(1,), dtype=int64
    [3]

Outputs:
  values: shape=(3, 3), dtype=float32
    [[ 3.,  2.,  1.],
     [ 7.,  6.,  5.],
     [11., 10.,  9.]]
  indices: shape=(3, 3), dtype=int64
    [[3, 2, 1],
     [3, 2, 1],
     [3, 2, 1]]

test_cc_top_k_neg_inf

Node:
  TopK(x, k) -> (values, indices)
  Attributes:
    axis = 1
    largest = 0
    sorted = 1

Inputs:
  x: shape=(1, 5), dtype=float32
    [[  1., -inf,  -3., -inf,   2.]]
  k: shape=(1,), dtype=int64
    [3]

Outputs:
  values: shape=(1, 3), dtype=float32
    [[-inf, -inf,  -3.]]
  indices: shape=(1, 3), dtype=int64
    [[1, 3, 2]]

test_cc_top_k_negative_axis

Node:
  TopK(x, k) -> (values, indices)
  Attributes:
    axis = -1

Inputs:
  x: shape=(3, 4), dtype=float32
    [[ 0.,  1.,  2.,  3.],
     [ 4.,  5.,  6.,  7.],
     [ 8.,  9., 10., 11.]]
  k: shape=(1,), dtype=int64
    [3]

Outputs:
  values: shape=(3, 3), dtype=float32
    [[ 3.,  2.,  1.],
     [ 7.,  6.,  5.],
     [11., 10.,  9.]]
  indices: shape=(3, 3), dtype=int64
    [[3, 2, 1],
     [3, 2, 1],
     [3, 2, 1]]

test_cc_top_k_pos_inf

Node:
  TopK(x, k) -> (values, indices)
  Attributes:
    axis = 1
    largest = 1
    sorted = 1

Inputs:
  x: shape=(1, 5), dtype=float32
    [[ 1., inf,  3., inf,  2.]]
  k: shape=(1,), dtype=int64
    [3]

Outputs:
  values: shape=(1, 3), dtype=float32
    [[inf, inf,  3.]]
  indices: shape=(1, 3), dtype=int64
    [[1, 3, 2]]

test_cc_top_k_pos_neg_inf

Node:
  TopK(x, k) -> (values, indices)
  Attributes:
    axis = 0
    largest = 1
    sorted = 1

Inputs:
  x: shape=(5,), dtype=float32
    [ inf,  -1., -inf,   2.,   0.]
  k: shape=(1,), dtype=int64
    [3]

Outputs:
  values: shape=(3,), dtype=float32
    [inf,  2.,  0.]
  indices: shape=(3,), dtype=int64
    [0, 3, 4]

test_cc_top_k_same_values

Node:
  TopK(x, k) -> (values, indices)
  Attributes:
    axis = 0

Inputs:
  x: shape=(5,), dtype=float32
    [1., 2., 3., 3., 2.]
  k: shape=(1,), dtype=int64
    [3]

Outputs:
  values: shape=(3,), dtype=float32
    [3., 3., 2.]
  indices: shape=(3,), dtype=int64
    [2, 3, 1]

test_cc_top_k_same_values_2d

Node:
  TopK(x, k) -> (values, indices)
  Attributes:
    axis = 1

Inputs:
  x: shape=(2, 4), dtype=float32
    [[1., 2., 2., 3.],
     [5., 5., 4., 3.]]
  k: shape=(1,), dtype=int64
    [3]

Outputs:
  values: shape=(2, 3), dtype=float32
    [[3., 2., 2.],
     [5., 5., 4.]]
  indices: shape=(2, 3), dtype=int64
    [[3, 1, 2],
     [0, 1, 2]]

test_cc_top_k_same_values_largest

Node:
  TopK(x, k) -> (values, indices)
  Attributes:
    axis = 0
    largest = 1
    sorted = 1

Inputs:
  x: shape=(5,), dtype=float32
    [1., 2., 3., 3., 2.]
  k: shape=(1,), dtype=int64
    [3]

Outputs:
  values: shape=(3,), dtype=float32
    [3., 3., 2.]
  indices: shape=(3,), dtype=int64
    [2, 3, 1]

test_cc_top_k_smallest

Node:
  TopK(x, k) -> (values, indices)
  Attributes:
    axis = 1
    largest = 0
    sorted = 1

Inputs:
  x: shape=(3, 4), dtype=float32
    [[ 0.,  1.,  2.,  3.],
     [ 4.,  5.,  6.,  7.],
     [11., 10.,  9.,  8.]]
  k: shape=(1,), dtype=int64
    [3]

Outputs:
  values: shape=(3, 3), dtype=float32
    [[ 0.,  1.,  2.],
     [ 4.,  5.,  6.],
     [ 8.,  9., 10.]]
  indices: shape=(3, 3), dtype=int64
    [[0, 1, 2],
     [0, 1, 2],
     [3, 2, 1]]

test_cc_top_k_uint64

Node:
  TopK(x, k) -> (values, indices)
  Attributes:
    axis = 1

Inputs:
  x: shape=(3, 4), dtype=uint64
    [[ 0,  1,  2,  3],
     [ 4,  5,  6,  7],
     [ 8,  9, 10, 11]]
  k: shape=(1,), dtype=int64
    [3]

Outputs:
  values: shape=(3, 3), dtype=uint64
    [[ 3,  2,  1],
     [ 7,  6,  5],
     [11, 10,  9]]
  indices: shape=(3, 3), dtype=int64
    [[3, 2, 1],
     [3, 2, 1],
     [3, 2, 1]]

Differences with previous version (10)#

SchemaDiff: TopK (domain 'ai.onnx')

  • old version: 10

  • new version: 11

  • breaking: no

Attributes:

  • added ‘largest’: type=INT; required=False; default=1

  • added ‘sorted’: type=INT; required=False; default=1

Type constraints:

  • changed ‘T’: added types: [‘tensor(int16)’, ‘tensor(int32)’, ‘tensor(int64)’, ‘tensor(int8)’, ‘tensor(uint16)’, ‘tensor(uint32)’, ‘tensor(uint64)’, ‘tensor(uint8)’]

Documentation:

  • line similarity: 0.22 (+13/-8 lines)

--- TopK v10
+++ TopK v11
@@ -1,11 +1,16 @@

-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
 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.
+* 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) then the k largest elements are returned.
+* If "sorted" is 1 (the default) 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.

Version History#