文件最后提交记录最后更新时间
1 个月前
2 个月前
2 个月前
2 个月前
2 个月前
README

RadixTopK

产品支持情况

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

功能说明

  • 接口功能:返回输入Tensor在指定维度上的k个极值及索引。RadixTopK算子支持通过输入Tensor指定k值。

  • 计算说明:

    • 当largest=True时,返回指定维度上最大的k个值及其索引
    • 当largest=False时,返回指定维度上最小的k个值及其索引

参数说明

参数名 输入/输出/属性 描述 数据类型 数据格式
x 输入 待进行TopK计算的输入张量,支持1-8维度。 FLOAT16、BFLOAT16 ND
k 输入 0D标量Tensor,表示计算维度上输出的极值个数。取值范围为[0, x.size(dim)]。 INT32、INT64 ND
values 输出 TopK计算的输出值,数据类型与x保持一致。shape排序轴与k一致,非排序轴与x一致。 FLOAT16、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 -

约束说明