Gather -> ShapeOf simplification (#24954)

### Details:
 - *Gather -> ShapeOf simplification*

### Tickets:
 - *CVS-143648*
This commit is contained in:
Evgenya Nugmanova 2024-06-12 10:53:05 +04:00 committed by GitHub
parent 43de73cbd6
commit c1ace6c0cb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 7 deletions

View File

@ -178,21 +178,21 @@ pass::AbsSinking::AbsSinking() {
pass::SimplifyGatherShapeOf::SimplifyGatherShapeOf() {
MATCHER_SCOPE(SimplifyGatherShapeOf);
const auto gather_pattern = wrap_type<op::util::GatherBase>();
const auto data_pattern = pattern::any_input(pattern::has_static_rank());
const auto indices_pattern = pattern::any_input(pattern::has_static_rank());
const auto axis_pattern = wrap_type<op::v0::Constant>();
const auto gather_pattern = wrap_type<op::util::GatherBase>({data_pattern, indices_pattern, axis_pattern});
const auto shape_of_pattern = wrap_type<v0::ShapeOf, v3::ShapeOf>({gather_pattern});
matcher_pass_callback callback = [](Matcher& m) {
auto node = m.get_match_root();
auto gather = as_type_ptr<v1::Gather>(node->input_value(0).get_node_shared_ptr());
if (!gather) {
auto gather = as_type_ptr<op::util::GatherBase>(node->input_value(0).get_node_shared_ptr());
if (!gather || gather->get_batch_dims() != 0) {
return false;
}
int64_t axis = gather->get_axis();
auto gather_in_rank = gather->get_input_partial_shape(0).rank();
auto indices_rank = gather->get_input_partial_shape(1).rank();
auto axis = gather->get_axis();
if (gather_in_rank.is_dynamic() || indices_rank.is_dynamic() || axis == v1::Gather::AXIS_NOT_SET_VALUE) {
return false;
}
auto zero_axis = v0::Constant::create<int64_t>(element::i64, Shape{}, {0});
NodeVector new_ops;