Tasks

mint.topk 算子

Background

1. 标杆情况

• 标杆接口链接：
  https://pytorch.org/docs/stable/generated/torch.topk.html#torch-topk
  [torch.topk]

• 标杆支持数据类型：xxx

2. MindSpore算子情况

• 当前支持数据类型

  CPU：不转测
  GPU：不转测
  Ascend：INT8、UINT8、INT16、INT32、INT64、FLOAT16、FLOAT32、DOUBLE、BFLOAT16(910B)
  

• 三后端统一后算子支持（标杆支持+三后端并集）
  INT8、UINT8、INT16、INT32、INT64、FLOAT16、FLOAT32、DOUBLE、BFLOAT16(910B)

Introduction

1. 功能介绍

沿给定维度查找 k 个最大或最小元素和对应的索引。

2. 接口描述（nn/functional接口需要与标杆保持一致）

• functional接口
  mindspore/python/mindspore/ops/function/array_func.py
  mindspore/python/mindspore/mint/init.py

  def topk_ext(input, k, dim=-1, largest=True, sorted=True): -> (tensor, tensor):
      return _get_cache_prim(ops.auto_generate.TopkExt)()(input, k, dim, largest, sorted)
  
  topk_ext -> mint.topk
  

• 算子原语
  mindspore/ops/auto_generate/gen_ops_prim.py
  class TopkExt(Primitive):

  参数	类型	输入/输出/属性	说明
  input	Tensor	输入
  k	int64	输入
  dim	int64	输入
  largest	bool	输入
  sorted	bool	输入
  output	Tensor	输出

Tasks
PyBoost接口
动态Shape
静态图KernelByKernel
ACLNN算子：
aclnnTopk

#I6FYA2:众筹需求模板：【众筹】存量算子 xxx 功能补齐（如数据类型补齐）及性能优化:众筹需求模板：【众筹】存量算子 xxx 功能补齐（如数据类型补齐）及性能优化