.. _op_ai_onnx_TopK: 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** .. code-block:: text Node: TopK(x, k) -> (values, indices) Attributes: axis = 1 .. code-block:: text 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** .. code-block:: text Node: TopK(x, k) -> (values, indices) Attributes: axis = 1 largest = 0 sorted = 1 .. code-block:: text 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** .. code-block:: text Node: TopK(x, k) -> (values, indices) Attributes: axis = -1 .. code-block:: text 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** .. code-block:: text Node: TopK(x, k) -> (values, indices) Attributes: axis = 1 largest = 1 sorted = 1 .. code-block:: text 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** .. code-block:: text Node: TopK(x, k) -> (values, indices) Attributes: axis = 0 largest = 1 sorted = 1 .. code-block:: text 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** .. code-block:: text Node: TopK(x, k) -> (values, indices) Attributes: axis = 0 .. code-block:: text 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** .. code-block:: text Node: TopK(x, k) -> (values, indices) Attributes: axis = 1 .. code-block:: text 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** .. code-block:: text Node: TopK(x, k) -> (values, indices) Attributes: axis = 0 largest = 1 sorted = 1 .. code-block:: text 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** .. code-block:: text Node: TopK(x, k) -> (values, indices) Attributes: axis = 1 largest = 0 sorted = 1 .. code-block:: text 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** .. code-block:: text Node: TopK(x, k) -> (values, indices) Attributes: axis = 1 .. code-block:: text 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) .. code-block:: diff --- 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 --------------- - :doc:`Version 10 ` - :doc:`Version 1 `