xxq250
/
mindspore
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

!17078 change the sequence of register info of topk. 

From: @wangshuide2020
Reviewed-by: @liangchenghui,@jjfeing
Signed-off-by: @liangchenghui
This commit is contained in:
mindspore-ci-bot 2021-05-27 15:31:51 +08:00 committed by Gitee
parent 61344a0600 54e50b40ee
commit a9efac9d32
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
mindspore/ops/_op_impl/tbe/top_k.py
View File

@ -24,14 +24,14 @@ top_k_op_info = TBERegOp("TopK") \
.kernel_name("top_k_d") \
.partial_flag(True) \
.need_check_supported(True) \
.attr("dim", "optional", "int", "all", "-1") \
.attr("k", "required", "int", "all") \
.attr("largest", "optional", "bool", "all", "true") \
.attr("sorted", "optional", "bool", "all", "true")\
.input(0, "x", False, "required", "all") \
.input(1, "assist_seq", False, "required", "all") \
.output(0, "values", False, "required", "all") \
.output(1, "indices", False, "required", "all") \
.attr("k", "required", "int", "all") \
.attr("sorted", "optional", "bool", "all", "true")\
.attr("dim", "optional", "int", "all", "-1") \
.attr("largest", "optional", "bool", "all", "true") \
.dtype_format(DataType.F16_Default, DataType.F16_Default, DataType.F16_Default, DataType.I32_Default) \
.get_op_info()