IntenseWebs/openvino
3
0 Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[PT FE]: support aten::broadcast_tensors (#19994)

Browse Source 
* broadcast tensors

* [PT FE]: support aten::broadcast_tensors

* apply review comments

* remove add
This commit is contained in:
Ekaterina Aidova authored and GitHub committed 2023-09-22 13:54:44 +04:00
1 parent 2151e5f979
commit 26d18c924b
2 files changed
+68
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/frontends/pytorch/src/transforms/prim_list_unpack_replacer.cpp
+23
View File
@@ -173,6 +173,29 @@ PrimListUnpackReplacer::PrimListUnpackReplacer() {
}
}
if (auto broadcast_tensors = cast_fw_node(input_node, "aten::broadcast_tensors")) {
auto tensors = cast_fw_node(broadcast_tensors->input_value(0).get_node_shared_ptr(), "prim::ListConstruct");
if (!tensors) {
add_exception_to_fw_node(input_node,
"aten::broadcast_tensors: only prim::ListConstruct supported as input.");
return false;
}
Output<Node> final_shape_t = opset10::Constant::create(element::i32, Shape{}, {0});
for (auto input : tensors->inputs()) {
auto tensor_shape = rg.make<opset10::ShapeOf>(input.get_source_output(), element::i32);
final_shape_t =
rg.make<opset10::Broadcast>(final_shape_t, tensor_shape, ov::op::BroadcastType::BIDIRECTIONAL);
}
auto final_shape = rg.make<opset10::ShapeOf>(final_shape_t, element::i32);
OutputVector outputs;
for (auto input : tensors->inputs()) {
outputs.push_back(rg.make<opset10::Broadcast>(input.get_source_output(), final_shape));
}
copy_runtime_info_and_name(list_unpack, rg.get(), {input_node});
replace_node(list_unpack, outputs);
return true;
}
if (auto unbind = cast_fw_node(input_node, "aten::unbind")) {
const auto input = unbind->get_input_source_output(0);
const auto axis = unbind->get_input_source_output(1);