pj3fizlut
/
mindspore
forked from OSSInnovation/mindspore
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

!1582 add topk and randomchoicewithmask op data type for aicpu 

Merge pull request !1582 from yanzhenxiang2020/r03_add_datatype
This commit is contained in:
mindspore-ci-bot 2020-05-28 19:43:41 +08:00 committed by Gitee
parent 02914ba0b9 d5af2f23b2
commit b70b2da675
2 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
mindspore/ops/_op_impl/aicpu/random_choice_with_mask.py
View File

@ -25,6 +25,7 @@ random_choice_with_mask_op_info = AiCPURegOp("RandomChoiceWithMask") \
.attr("seed", "int") \
.attr("seed2", "int") \
.dtype_format(DataType.BOOL_NCHW, DataType.I32_NCHW, DataType.BOOL_NCHW) \
.dtype_format(DataType.BOOL_Default, DataType.I32_Default, DataType.BOOL_Default) \
.get_op_info()
@op_info_register(random_choice_with_mask_op_info)

1
mindspore/ops/_op_impl/aicpu/topk.py
View File

@ -24,6 +24,7 @@ top_k_op_info = AiCPURegOp("TopK") \
.output(0, "values", "required") \
.output(1, "indices", "required") \
.dtype_format(DataType.F16_Default, DataType.I32_Default, DataType.F16_Default, DataType.I32_Default) \
.dtype_format(DataType.I32_Default, DataType.I32_Default, DataType.I32_Default, DataType.I32_Default) \
.get_op_info()
@op_info_register(top_k_op_info)