Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[PT FE] Return default NonZero type where possible #29072

Merged
merged 2 commits into from
Feb 20, 2025
Merged

[PT FE] Return default NonZero type where possible #29072

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f4b6585
[PT FE] Return default NonZero type where possible
mvafin Feb 19, 2025
2ebfc84
Merge branch 'master' into mvafin/pt_fe/nonzero_def
mvafin Feb 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/frontends/pytorch/src/op/index.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ OutputVector translate_index(const NodeContext& context) {
}
}
if (index_ov_type == element::boolean || index_ov_type == element::u8) {
auto nonzero = context.mark_node(std::make_shared<v3::NonZero>(indices, element::i32));
auto nonzero = context.mark_node(std::make_shared<v3::NonZero>(indices));
auto input_order = context.mark_node(v0::Constant::create(element::i32, Shape{2}, {1, 0}));
auto masked_id = context.mark_node(std::make_shared<v1::Transpose>(nonzero, input_order));
auto gather = context.mark_node(std::make_shared<v8::GatherND>(x, masked_id));
Expand Down
2 changes: 1 addition & 1 deletion src/frontends/pytorch/src/op/masked_scatter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ OutputVector translate_masked_scatter(const NodeContext& context) {
auto x_shape = context.mark_node(std::make_shared<v3::ShapeOf>(x, element::i32));
auto expanded_mask =
context.mark_node(std::make_shared<v3::Broadcast>(mask, x_shape, BroadcastType::BIDIRECTIONAL));
auto index = context.mark_node(std::make_shared<v3::NonZero>(expanded_mask, element::i32));
auto index = context.mark_node(std::make_shared<v3::NonZero>(expanded_mask));
auto input_order = context.mark_node(v0::Constant::create(element::i32, Shape{2}, {1, 0}));
index = context.mark_node(std::make_shared<v1::Transpose>(index, input_order));
// source can be arbitary shape, select only relevant data
Expand Down
2 changes: 1 addition & 1 deletion src/frontends/pytorch/src/transforms/aten_index_put_replacer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,7 @@ AtenIndexPutReplacer::AtenIndexPutReplacer() {
auto input_shape = rg.make<v3::ShapeOf>(input, element::i32);
auto input_rank = rg.make<v3::ShapeOf>(input_shape, element::i32);
auto one_const = v0::Constant::create(element::i32, Shape{1}, {1});
auto nonzero = rg.make<v3::NonZero>(index, element::i32);
auto nonzero = rg.make<v3::NonZero>(index);
auto input_order = v0::Constant::create(element::i32, Shape{2}, {1, 0});
index = rg.make<v1::Transpose>(nonzero, input_order);
auto result = rg.make<v3::ScatterNDUpdate>(input, index, values);
Expand Down
2 changes: 1 addition & 1 deletion src/frontends/pytorch/src/transforms/aten_index_replacer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ AtenIndexToSelect::AtenIndexToSelect() {
}
auto index_dtype = indicies->get_output_element_type(0);
if (index_dtype == element::boolean || index_dtype == element::u8) {
auto nonzero = rg.make<v3::NonZero>(indicies, element::i32);
auto nonzero = rg.make<v3::NonZero>(indicies);
auto input_order = v0::Constant::create(element::i32, Shape{2}, {1, 0});
auto masked_id = rg.make<v1::Transpose>(nonzero, input_order);
auto gather = rg.make<v8::GatherND>(input_node, masked_id);
Expand Down
2 changes: 1 addition & 1 deletion src/frontends/pytorch/src/utils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -728,7 +728,7 @@ bool index_tensor_on_list(ov::pass::NodeRegistry& rg,
auto id_dtype = indices[i].get_element_type();
if (id_dtype == element::boolean || id_dtype == element::u8) {
auto idx = rg.make<v0::Convert>(indices[i], element::u8);
auto nonzero = rg.make<v3::NonZero>(idx, element::i32);
auto nonzero = rg.make<v3::NonZero>(idx);
auto input_order = rg.make<v0::Constant>(element::i32, Shape{2}, std::vector<int32_t>{1, 0});
auto masked_id = rg.make<v1::Transpose>(nonzero, input_order);
masked_indicies.push_back(masked_id);
Expand Down
Loading