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). Default value is -1.

• largest: Whether to return the top-K largest or smallest elements. Default value is 1.

• sorted: Whether to return the elements in sorted order. Default value is 1.

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

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

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

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')

Differences

0	0	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
1	1	shape [a_1, a_2, ..., a_n, r] and integer argument k, return two outputs:	shape [a_1, a_2, ..., a_n, r] and integer argument k, return two outputs:
2	2	-Value tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]	-Value tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]
3	3	which contains the values of the top k elements along the specified axis	which contains the values of the top k elements along the specified axis
4	4	-Index tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] which	-Index tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] which
5	5	contains the indices of the top k elements (original indices from the input	contains the indices of the top k elements (original indices from the input
6	6	tensor).	tensor).
7	7
	8		If "largest" is 1 (the default value) then the k largest elements are returned.
	9		If "sorted" is 1 (the default value) then the resulting k elements will be sorted.
	10		If "sorted" is 0, order of returned 'Values' and 'Indices' are undefined.
	11
8	12	Given two equivalent values, this operator uses the indices along the axis as	Given two equivalent values, this operator uses the indices along the axis as
9	13	a tiebreaker. That is, the element with the lower index will appear first.	a tiebreaker. That is, the element with the lower index will appear first.
10	14
11	15	**Attributes**	**Attributes**
12	16
13	17	* **axis**:	* **axis**:
14	18	Dimension on which to do the sort. Default value is -1.	Dimension on which to do the sort. Negative value means counting
	19		dimensions from the back. Accepted range is [-r, r-1] where r =
	20		rank(input). Default value is -1.
	21		* **largest**:
	22		Whether to return the top-K largest or smallest elements. Default value is 1.
	23		* **sorted**:
	24		Whether to return the elements in sorted order. Default value is 1.
15	25
16	26	**Inputs**	**Inputs**
17	27
18	28	* **X** (heterogeneous) - **T**:	* **X** (heterogeneous) - **T**:
19	29	Tensor of shape [a_1, a_2, ..., a_n, r]	Tensor of shape [a_1, a_2, ..., a_n, r]
20	30	* **K** (heterogeneous) - **tensor(int64)**:	* **K** (heterogeneous) - **tensor(int64)**:
21	31	A 1-D tensor containing a single positive value corresponding to the	A 1-D tensor containing a single positive value corresponding to the
22	32	number of top elements to retrieve	number of top elements to retrieve
23	33
24	34	**Outputs**	**Outputs**
25	35
26	36	* **Values** (heterogeneous) - **T**:	* **Values** (heterogeneous) - **T**:
27	37	Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]	Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]
28	38	containing top K values from the input tensor	containing top K values from the input tensor
29	39	* **Indices** (heterogeneous) - **I**:	* **Indices** (heterogeneous) - **I**:
30	40	Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]	Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]
31	41	containing the corresponding input tensor indices for the top K	containing the corresponding input tensor indices for the top K
32	42	values.	values.
33	43
34	44	**Type Constraints**	**Type Constraints**
35	45
36	46	* **T** in (	* **T** in (
37	47	tensor(double),	tensor(double),
38	48	tensor(float),	tensor(float),
39	49	tensor(float16)	tensor(float16),
	50		tensor(int16),
	51		tensor(int32),
	52		tensor(int64),
	53		tensor(int8),
	54		tensor(uint16),
	55		tensor(uint32),
	56		tensor(uint64),
	57		tensor(uint8)
40	58	):	):
41	59	Constrain input and output types to float tensors.	Constrain input and output types to numeric tensors.
42	60	* **I** in (	* **I** in (
43	61	tensor(int64)	tensor(int64)
44	62	):	):
45	63	Constrain index tensor to 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_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. Default value is -1.

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

Differences

0	0	Retrieve the top-K elements along a specified axis. Given an input tensor of	Retrieve the top-K elements along a specified axis. Given an input tensor of
1	1	shape [a_1, a_2, ..., a_n, r] and integer argument k, return two outputs:	shape [a_1, a_2, ..., a_n, r] and integer argument k, return two outputs:
2	2	-Value tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]	-Value tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]
3	3	which contains the values of the top k elements along the specified axis	which contains the values of the top k elements along the specified axis
4	4	-Index tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] which	-Index tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n] which
5	5	contains the indices of the top k elements (original indices from the input	contains the indices of the top k elements (original indices from the input
6	6	tensor).	tensor).
	7
7	8	Given two equivalent values, this operator uses the indices along the axis as	Given two equivalent values, this operator uses the indices along the axis as
8	9	a tiebreaker. That is, the element with the lower index will appear first.	a tiebreaker. That is, the element with the lower index will appear first.
9	10
10	11	**Attributes**	**Attributes**
11	12
12	13	* **axis**:	* **axis**:
13	14	Dimension on which to do the sort. Default value is -1.	Dimension on which to do the sort. Default value is -1.
14		* **k** (required):
15		Number of top elements to retrieve
16	15
17	16	**Inputs**	**Inputs**
18	17
19	18	* **X** (heterogeneous) - **T**:	* **X** (heterogeneous) - **T**:
20	19	Tensor of shape [a_1, a_2, ..., a_n, r]	Tensor of shape [a_1, a_2, ..., a_n, r]
	20		* **K** (heterogeneous) - **tensor(int64)**:
	21		A 1-D tensor containing a single positive value corresponding to the
	22		number of top elements to retrieve
21	23
22	24	**Outputs**	**Outputs**
23	25
24	26	* **Values** (heterogeneous) - **T**:	* **Values** (heterogeneous) - **T**:
25	27	Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]	Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]
26	28	containing top K values from the input tensor	containing top K values from the input tensor
27	29	* **Indices** (heterogeneous) - **I**:	* **Indices** (heterogeneous) - **I**:
28	30	Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]	Tensor of shape [a_1, a_2, ..., a_{axis-1}, k, a_{axis+1}, ... a_n]
29	31	containing the corresponding input tensor indices for the top K	containing the corresponding input tensor indices for the top K
30	32	values.	values.
31	33
32	34	**Type Constraints**	**Type Constraints**
33	35
34	36	* **T** in (	* **T** in (
35	37	tensor(double),	tensor(double),
36	38	tensor(float),	tensor(float),
37	39	tensor(float16)	tensor(float16)
38	40	):	):
39	41	Constrain input and output types to float tensors.	Constrain input and output types to float tensors.
40	42	* **I** in (	* **I** in (
41	43	tensor(int64)	tensor(int64)
42	44	):	):
43	45	Constrain index tensor to 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. Default value is -1.

• 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

previous

Tile

next

Transpose