TopK

English | 简体中文

Product Support Status

Product	Supported
Ascend 950PR&950DT products	√
Atlas A3 products	√
Atlas A2 products	√
Atlas 200I/500 A2 inference products	√
Atlas inference products	√
Atlas training products	√

Function Description

• Operator function: finds the k largest or smallest elements in the input tensor and their indices along the specified dimension.

Parameter Description

Parameter Name	Input/Output/Attribute	Description	Data Type	Data Format
x	Input	Input tensor.	FLOAT16, FLOAT, DOUBLE, INT8, INT16, INT32, INT64, UINT8, UINT16, UINT32, UINT64, BFLOAT16	ND
k	Input	Number of elements to take.	INT32	ND
sorted	Attribute	Whether to sort the output results. The default value is true.	Bool	-
largest	Attribute	Whether to take the maximum values. The default value is true. If false, the minimum values are taken.	Bool	-
dim	Attribute	Dimension along which the operation is performed. The default value is -1 (the last dimension).	Int	-
values	Output	Output top-k element values.	FLOAT16, FLOAT, DOUBLE, INT8, INT16, INT32, INT64, UINT8, UINT16, UINT32, UINT64, BFLOAT16	ND
indices	Output	Indices of the output top-k elements in the input tensor.	INT32	ND

Constraint Description

• dim must be in the range [-number of dimensions of x, number of dimensions of x).
• k must be greater than or equal to 0 and less than or equal to the size of x on the dim dimension.
• The BFLOAT16 type is supported only on specific products.

Invocation Description

Invocation Method	Sample Code	Description
Graph-mode invocation	test_geir_topk	Invokes the TopK operator through the operator IR API.