提交 d5af2f23 编写于 作者: Y yanzhenxiang2020

add topk and randomchoicewithmask data type for aicpu

上级 431bc8bf
......@@ -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)
......
......@@ -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)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册