TopK

产品支持情况

产品	是否支持
Ascend 950PR/Ascend 950DT	√
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas A2 训练系列产品/Atlas A2 推理系列产品	√
Atlas 200I/500 A2 推理产品	√
Atlas 推理系列产品	√
Atlas 训练系列产品	√

功能说明

• 算子功能：沿指定维度找出输入张量中最大或最小的k个元素及其索引。

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
x	输入	输入张量。	FLOAT16、FLOAT、DOUBLE、INT8、INT16、INT32、INT64、UINT8、UINT16、UINT32、UINT64、BFLOAT16	ND
k	输入	要取出的元素个数。	INT32	ND
sorted	属性	是否对输出结果排序，默认为true。	Bool	-
largest	属性	是否取最大值，默认为true。若为false则取最小值。	Bool	-
dim	属性	指定沿哪个维度进行操作，默认为-1（最后一维）。	Int	-
values	输出	输出的top-k个元素值。	FLOAT16、FLOAT、DOUBLE、INT8、INT16、INT32、INT64、UINT8、UINT16、UINT32、UINT64、BFLOAT16	ND
indices	输出	输出的top-k个元素在输入张量中的索引。	INT32	ND

约束说明

• dim必须在[-x维度数, x维度数)范围内
• k必须大于等于0且小于等于x在dim维度上的大小
• BFLOAT16类型仅在特定产品上支持

调用说明

调用方式	样例代码	说明
图模式接口	test_geir_topk	通过算子IR接口方式调用TopK算子。