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

change the sequence of register info of topk.

This commit is contained in:
wangshuide2020 2021-05-27 10:44:04 +08:00
parent 1691ab80de
commit 54e50b40ee
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()