text: // Copyright (C) 2018-2021 Intel Corporation
  // SPDX-License-Identifier: Apache-2.0
  //
  
  #include "legacy/transformations/convert_opset1_to_legacy/convert_cells_to_cells_ie.hpp"
  
  #include <memory>
  #include <vector>
  
  #include <ngraph/opsets/opset1.hpp>
  #include <ngraph/opsets/opset3.hpp>
  #include <ngraph/opsets/opset4.hpp>
  #include <ngraph/rt_info.hpp>
  #include <ngraph/pattern/op/wrap_type.hpp>
  #include <ngraph/op/util/rnn_cell_base.hpp>
  
  #include "legacy/ngraph_ops/lstm_cell_ie.hpp"
  #include "legacy/ngraph_ops/gru_cell_ie.hpp"
  #include "legacy/ngraph_ops/rnn_cell_ie.hpp"
  
  NGRAPH_RTTI_DEFINITION(ngraph::pass::ConvertLSTMCellMatcher, "ConvertLSTMCellMatcher", 0);
  
  ngraph::pass::ConvertLSTMCellMatcher::ConvertLSTMCellMatcher() {
      auto is_supported_lstm_cell = [](const std::shared_ptr<Node>& n) {
          return pattern::has_class<ngraph::opset1::LSTMCell>()(n) || pattern::has_class<ngraph::opset4::LSTMCell>()(n);
      };
      auto any_lstm = std::make_shared<pattern::op::Label>(element::f32, Shape{}, is_supported_lstm_cell);
      ngraph::matcher_pass_callback callback = [](pattern::Matcher& m) {
          auto lstm_cell = std::dynamic_pointer_cast<ngraph::op::util::RNNCellBase>(m.get_match_root());
          if (!lstm_cell) {
              return false;
          }
          auto W = std::dynamic_pointer_cast<ngraph::opset1::Constant> (lstm_cell->input_value(3).get_node_shared_ptr());
          if (!W) {
              return false;
          }
  
          auto R = std::dynamic_pointer_cast<ngraph::opset1::Constant> (lstm_cell->input_value(4).get_node_shared_ptr());
          if (!R) {
              return false;
          }
  
          auto concat = std::make_shared<ngraph::opset1::Concat>(ngraph::NodeVector({W, R}), 1);
          auto lstm_cell_ie = std::make_shared<ngraph::op::LSTMCellIE> (lstm_cell->input(0).get_source_output(),  // X
                                                                        lstm_cell->input(1).get_source_output(),  // initial_hidden_state
                                                                        lstm_cell->input(2).get_source_output(),  // initial_cell_state
                                                                        concat->output(0),                       // WR
                                                                        lstm_cell->input(5).get_source_output(),  // B
                                                                        lstm_cell->get_hidden_size(),
                                                                        lstm_cell->get_activations(),
                                                                        lstm_cell->get_activations_alpha(),
                                                                        lstm_cell->get_activations_beta(),
                                                                        lstm_cell->get_clip());
  
          lstm_cell_ie->set_friendly_name(lstm_cell->get_friendly_name());
          ngraph::copy_runtime_info(lstm_cell, {concat, lstm_cell_ie});
          ngraph::replace_node(m.get_match_root(), lstm_cell_ie);
          return true;
      };
  
      auto m = std::make_shared<ngraph::pattern::Matcher>(any_lstm, "ConvertLSTMCellToLSTMCellIE");
      this->register_matcher(m, callback);
  }
  
  NGRAPH_RTTI_DEFINITION(ngraph::pass::ConvertGRUCellMatcher, "ConvertGRUCellMatcher", 0);
  
  ngraph::pass::ConvertGRUCellMatcher::ConvertGRUCellMatcher() {
      auto gru_cell_ngraph = ngraph::pattern::wrap_type<ngraph::opset3::GRUCell>();
  
      ngraph::matcher_pass_callback callback = [](pattern::Matcher& m) {
          auto gru_cell = std::dynamic_pointer_cast<ngraph::opset3::GRUCell> (m.get_match_root());
          if (!gru_cell) {
              return false;
          }
  
          auto W = std::dynamic_pointer_cast<ngraph::opset1::Constant> (gru_cell->input_value(2).get_node_shared_ptr());
          if (!W) {
              return false;
          }
  
          auto R = std::dynamic_pointer_cast<ngraph::opset1::Constant> (gru_cell->input_value(3).get_node_shared_ptr());
          if (!R) {
              return false;
          }
  
          auto concat = std::make_shared<ngraph::opset1::Concat>(ngraph::NodeVector({W, R}), 1);
          auto gru_cell_ie = std::make_shared<ngraph::op::GRUCellIE>(gru_cell->input(0).get_source_output(),  // X
                                                                     gru_cell->input(1).get_source_output(),  // initial_hidden_state
                                                                     concat->output(0),                       // WR
                                                                     gru_cell->input(4).get_source_output(),  // B
                                                                     gru_cell->get_hidden_size(),
                                                                     gru_cell->get_activations(),
                                                                     gru_cell->get_activations_alpha(),
                                                                     gru_cell->get_activations_beta(),
                                                                     gru_cell->get_clip(),
                                                                     gru_cell->get_linear_before_reset());
  
          gru_cell_ie->set_friendly_name(gru_cell->get_friendly_name());
          ngraph::copy_runtime_info(gru_cell, {concat, gru_cell_ie});
          ngraph::replace_node(m.get_match_root(), gru_cell_ie);
          return true;
      };
  
      auto m = std::make_shared<ngraph::pattern::Matcher>(gru_cell_ngraph, "ConvertGRUCellToGRUCellIE");
      this->register_matcher(m, callback);
  }
  
  NGRAPH_RTTI_DEFINITION(ngraph::pass::ConvertRNNCellMatcher, "ConvertRNNCellMatcher", 0);
  
  ngraph::pass::ConvertRNNCellMatcher::ConvertRNNCellMatcher() {
      auto rnn_cell_ngraph = ngraph::pattern::wrap_type<ngraph::opset3::RNNCell>();
  
      ngraph::matcher_pass_callback callback = [](pattern::Matcher& m) {
          auto rnn_cell = std::dynamic_pointer_cast<ngraph::opset3::RNNCell> (m.get_match_root());
          if (!rnn_cell) {
              return false;
          }
  
          auto W = std::dynamic_pointer_cast<ngraph::opset1::Constant> (rnn_cell->input_value(2).get_node_shared_ptr());
          if (!W) {
              return false;
          }
  
          auto R = std::dynamic_pointer_cast<ngraph::opset1::Constant> (rnn_cell->input_value(3).get_node_shared_ptr());
          if (!R) {
              return false;
          }
  
          auto concat = std::make_shared<ngraph::opset1::Concat>(ngraph::NodeVector({W, R}), 1);
          auto rnn_cell_ie = std::make_shared<ngraph::op::RNNCellIE> (rnn_cell->input(0).get_source_output(),  // X
                                                                      rnn_cell->input(1).get_source_output(),  // initial_hidden_state
                                                                      concat->output(0),                      // WR
                                                                      rnn_cell->input(4).get_source_output(),  // B
                                                                      rnn_cell->get_hidden_size(),
                                                                      rnn_cell->get_activations(),
                                                                      rnn_cell->get_activations_alpha(),
                                                                      rnn_cell->get_activations_beta(),
                                                                      rnn_cell->get_clip());
  
          rnn_cell_ie->set_friendly_name(rnn_cell->get_friendly_name());
          ngraph::copy_runtime_info(rnn_cell, {concat, rnn_cell_ie});
          ngraph::replace_node(m.get_match_root(), rnn_cell_ie);
          return true;
      };
  
      auto m = std::make_shared<ngraph::pattern::Matcher>(rnn_cell_ngraph, "ConvertRNNCellToRNNCellIE");
      this->register_matcher(m, callback);
  }
  
  