TopKV2

产品支持情况

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

功能说明

• 接口功能：返回输入Tensor在指定维度上的k个极值及索引。TopKV2算子支持通过输入Tensor指定k值，并支持可配置的indices输出数据类型。

• 计算说明：

  • 当largest=True时，返回指定维度上最大的k个值及其索引
  • 当largest=False时，返回指定维度上最小的k个值及其索引
  • 当sorted=True时，输出结果按照从大到小（largest=True）或从小到大（largest=False）排序
  • 当sorted=False时，输出结果不排序

参数说明

参数名	输入/输出/属性	描述	数据类型	数据格式
x	输入	待进行TopK计算的输入张量，支持1-8维度。	INT8、UINT8、INT16、UINT16、INT32、UINT32、INT64、UINT64、FLOAT16、FLOAT32、BFLOAT16	ND
k	输入	0D标量Tensor，表示计算维度上输出的极值个数。取值范围为[0, x.size(dim)]。	INT32、INT64	ND
values	输出	TopK计算的输出值，数据类型与x保持一致。shape排序轴与k一致，非排序轴与x一致。	INT8、UINT8、INT16、UINT16、INT32、UINT32、INT64、UINT64、FLOAT16、FLOAT32、BFLOAT16	ND
indices	输出	TopK计算的输出索引，数据类型由indices_dtype属性指定。shape排序轴与k一致，非排序轴与x一致。	INT32、INT64	ND
sorted	属性	可选布尔型，默认为True。True表示输出结果排序，False表示输出结果不排序。	Bool	-
dim	属性	可选整型，默认为-1。表示计算维度。取值范围为[-x.dim(), x.dim())。	Int	-
largest	属性	可选布尔型，默认为True。True表示返回最大的k个元素，False表示返回最小的k个元素。	Bool	-
indices_dtype	属性	可选整型，默认为DT_INT32(3)。表示输出indices的数据类型，支持DT_INT32(3)或DT_INT64(9)。	Int	-
sort_policy	属性	可选整型，默认为0。表示TopK候选结果的排序策略。0表示使用原有排序行为，1表示在sorted=True且小k值（2≤k≤32）场景下启用Bitonic Small TopK路径。	Int	-

• Atlas 200I/500 A2推理产品、Atlas推理系列产品、Atlas训练系列产品：数据类型不支持BFLOAT16。

约束说明

无

调用说明

调用方式	调用样例	说明
图模式调用	test_geir_top_k_v2	通过算子IR构图方式调用top_k_v2算子。