[PT FE] Support aten::pow_ (#21630)

* [PT FE] Support aten::pow_

* Fix tests
This commit is contained in:
Maxim Vafin
2023-12-13 18:54:53 +01:00
committed by GitHub
parent 336a69ec24
commit 4752237871
3 changed files with 38 additions and 19 deletions

View File

@@ -15,8 +15,17 @@ OutputVector translate_pow(const NodeContext& context) {
num_inputs_check(context, 2, 2);
auto lhs = context.get_input(0);
auto rhs = context.get_input(1);
align_eltwise_input_types(context, lhs, rhs, true);
return {context.mark_node(std::make_shared<ov::op::v1::Power>(lhs, rhs))};
auto inplace = context.get_op_type() == "aten::pow_";
if (inplace) {
rhs = std::make_shared<ov::op::v1::ConvertLike>(rhs, lhs);
} else {
align_eltwise_input_types(context, lhs, rhs, true);
}
auto res = context.mark_node(std::make_shared<ov::op::v1::Power>(lhs, rhs));
if (inplace) {
context.mutate_input(0, res);
}
return {res};
}
} // namespace op

View File

@@ -471,6 +471,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
{"aten::pixel_unshuffle", op::translate_pixel_unshuffle},
{"aten::prelu", op::translate_1to1_match_2_inputs<opset10::PRelu>},
{"aten::pow", op::translate_pow},
{"aten::pow_", op::translate_pow},
{"aten::prod", op::translate_prod},
{"aten::quantize_per_channel", op::translate_quantize_per_channel},
{"aten::quantize_per_tensor", op::translate_quantize_per_tensor},