| 文件 | 最后提交记录 | 最后更新时间 |
|---|---|---|
| 2 个月前 | ||
| 2 个月前 | ||
| 2 个月前 | ||
| 2 个月前 | ||
| 2 个月前 |
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 | - |
约束说明
无