(l-onnx-doc-TopK)= # TopK (l-onnx-op-topk-24)= ## TopK - 24 ### Version - **name**: [TopK (GitHub)](https://github.com/onnx/onnx/blob/main/docs/Operators.md#TopK) - **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 ```python 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 ```python 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 ```python 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 ```python 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 ```python 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 ```python 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 ```python 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", ) ``` ```{toctree} text_diff_TopK_11_24 ``` (l-onnx-op-topk-11)= ## TopK - 11 ### Version - **name**: [TopK (GitHub)](https://github.com/onnx/onnx/blob/main/docs/Operators.md#TopK) - **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 ```{toctree} text_diff_TopK_10_24 text_diff_TopK_10_11 ``` (l-onnx-op-topk-10)= ## TopK - 10 ### Version - **name**: [TopK (GitHub)](https://github.com/onnx/onnx/blob/main/docs/Operators.md#TopK) - **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 ```{toctree} text_diff_TopK_1_24 text_diff_TopK_1_11 text_diff_TopK_1_10 ``` (l-onnx-op-topk-1)= ## TopK - 1 ### Version - **name**: [TopK (GitHub)](https://github.com/onnx/onnx/blob/main/docs/Operators.md#TopK) - **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_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. * **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