--- /dev/null
+// Copyright (C) 2020 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "transformations/tensor_iterator_transformations/unroll_tensor_iterator.hpp"
+#include "transformations/utils/utils.hpp"
+
+#include <memory>
+#include <vector>
+
+#include <ngraph/graph_util.hpp>
+#include <ngraph/opsets/opset4.hpp>
+#include <ngraph/pattern/op/wrap_type.hpp>
+#include <ngraph/rt_info.hpp>
+
+ngraph::pass::UnrollTensorIterator::UnrollTensorIterator() : MatcherPass() {
+ auto tensor_iterator = ngraph::pattern::wrap_type<ngraph::opset4::TensorIterator>();
+ ngraph::matcher_pass_callback callback = [this](pattern::Matcher& m) {
+ auto ti = std::dynamic_pointer_cast<ngraph::opset4::TensorIterator>(m.get_match_root());
+ if (!ti) {
+ return false;
+ }
+
+ const auto function = ti->get_body()->to_function();
+ auto num_iter = ti->get_num_iterations();
+
+ // negative value means inconsistent TI
+ if (num_iter <= -1) {
+ return false;
+ }
+
+ // Create copies of the TensorIterator body, the number of copies is equal to the number of iterations.
+ // Assign names to the created layers.
+ std::vector<std::shared_ptr<ngraph::Function>> body_functions(num_iter);
+ for (uint64_t idx = 0; idx < num_iter; ++idx) {
+ body_functions[idx] = clone_function(*function);
+ for (auto &node : body_functions[idx]->get_ops()) {
+ node->set_friendly_name(ti->get_friendly_name() + "/" + std::to_string(idx + 1) + "/" + node->get_friendly_name());
+ copy_runtime_info(ti, node);
+ }
+ }
+
+ // Port map : inputs and back edges
+ for (const auto& desc : ti->get_input_descriptions()) {
+ const std::string& type_name = desc->get_type_info().name;
+
+ if (type_name == "SliceInputDescription") {
+ auto input_desc = std::dynamic_pointer_cast<ngraph::opset4::TensorIterator::SliceInputDescription>(desc);
+ if (!input_desc) {
+ return false;
+ }
+
+ // Connect the sliced input (layer before the input) to the Split layer and connect
+ // the corresponding Split output to the corresponding copy of the body.
+ // If the number of iterations is 1, then the Split is not needed.
+
+ auto in_data = ti->input_values()[input_desc->m_input_index];
+ const auto const_axis = opset4::Constant::create(element::i64, Shape{}, {input_desc->m_axis});
+
+ if (num_iter > 1) {
+ auto split = std::make_shared<ngraph::opset4::Split>(in_data, const_axis, num_iter);
+ copy_runtime_info(ti, split);
+ auto stride = input_desc->m_stride;
+ // connect to the body
+ for (uint64_t j = 0; j < num_iter; j++) {
+ auto idx = stride > 0 ? j : num_iter - j - 1;
+ auto param = body_functions[j]->get_parameters()[input_desc->m_body_parameter_index];
+ for (auto &output : param->outputs()) {
+ output.replace(split->output(idx));
+ }
+ }
+ } else {
+ // connect to the body
+ auto param = body_functions[0]->get_parameters()[input_desc->m_body_parameter_index];
+ for (auto &output : param->outputs()) {
+ output.replace(in_data);
+ }
+ }
+ } else if (type_name == "MergedInputDescription") {
+ auto input_desc = std::dynamic_pointer_cast<ngraph::opset4::TensorIterator::MergedInputDescription>(desc);
+ if (!input_desc) {
+ return false;
+ }
+
+ // Connect the input to the corresponding copy of the body.
+ auto in_data = ti->input_values()[input_desc->m_input_index].get_node_shared_ptr();
+ auto param = body_functions[0]->get_parameters()[input_desc->m_body_parameter_index];
+ for (auto &output : param->outputs()) {
+ output.replace(in_data);
+ }
+
+ // Back-edge processing. Connect the copies of the body to each other.
+ for (uint64_t j = 1; j < num_iter; j++) {
+ auto cur_param = body_functions[j]->get_parameters()[input_desc->m_body_parameter_index];
+ auto prev_val = body_functions[j - 1]->get_results()[input_desc->m_body_value_index];
+ for (auto &output : cur_param->outputs()) {
+ output.replace(prev_val->get_input_source_output(0));
+ }
+ }
+ } else if (type_name == "InvariantInputDescription") {
+ auto input_desc = std::dynamic_pointer_cast<ngraph::opset4::TensorIterator::InvariantInputDescription>(
+ desc);
+ if (!input_desc) {
+ return false;
+ }
+
+ // Connect the input to the corresponding copy of the body.
+ auto in_data = ti->input_values()[input_desc->m_input_index].get_node_shared_ptr();
+ for (uint64_t j = 0; j < num_iter; j++) {
+ auto param = body_functions[j]->get_parameters()[input_desc->m_body_parameter_index];
+ for (auto &output : param->outputs()) {
+ output.replace(in_data);
+ }
+ }
+ } else {
+ // "Incorrect type of the input description.";
+ return false;
+ }
+ }
+
+ // Port map: outputs
+ for (const auto& desc : ti->get_output_descriptions()) {
+ std::string type_name = desc->get_type_info().name;
+ if (type_name == "ConcatOutputDescription") {
+ auto output_desc = std::dynamic_pointer_cast<ngraph::opset4::TensorIterator::ConcatOutputDescription>(desc);
+ if (!output_desc) {
+ return false;
+ }
+
+ // Connect corresponding outputs (layers before Result op) of each copy of the body to Concat layer.
+ // Connect the Concat to corresponding output of TensorIterator.
+ // If the number of iterations is 1, then the Concat is not needed.
+
+ if (num_iter > 1) {
+ ngraph::OutputVector to_concat(num_iter);
+ auto stride = output_desc->m_stride;
+
+ // Connect outputs of the bodies to the Concat layer
+ for (uint64_t j = 0; j < num_iter; j++) {
+ auto idx = stride > 0 ? j : num_iter - j - 1;
+ std::shared_ptr<opset4::Result> result = body_functions[idx]->get_results()[output_desc->m_body_value_index];
+ auto input_to_res = result->get_input_source_output(0);
+ to_concat[j] = input_to_res;
+ }
+ auto concat = std::make_shared<ngraph::opset4::Concat>(to_concat, output_desc->m_axis);
+ copy_runtime_info(ti, concat);
+
+ // connect the Concat layer to the corresponding TI outputs
+ concat->output(0).get_tensor().set_name(
+ op::util::create_ie_output_name(ti->output(output_desc->m_output_index)));
+ for (auto &input : ti->output(output_desc->m_output_index).get_target_inputs()) {
+ input.replace_source_output(concat);
+ }
+ } else {
+ // Connect outputs of the bodies to the corresponding TI outputs
+ std::shared_ptr<opset4::Result> result = body_functions[0]->get_results()[output_desc->m_body_value_index];
+ auto input_to_res = result->get_input_source_output(0);
+ for (auto &input : ti->output(output_desc->m_output_index).get_target_inputs()) {
+ input.replace_source_output(input_to_res);
+ }
+ }
+ } else if (type_name == "BodyOutputDescription") {
+ auto output_desc = std::dynamic_pointer_cast<ngraph::opset4::TensorIterator::BodyOutputDescription>(desc);
+ if (!output_desc) {
+ return false;
+ }
+
+ // Connect outputs of the bodies to the corresponding TI outputs
+ auto iter = output_desc->m_iteration;
+ iter = iter >= 0? iter: num_iter - 1;
+ std::shared_ptr<opset4::Result> result = body_functions[iter]->get_results()[output_desc->m_body_value_index];
+ const auto& in_value = result->input_value(0);
+
+ in_value.get_tensor().set_name(op::util::create_ie_output_name(ti->output(output_desc->m_output_index)));
+ for (const auto &input : ti->output(output_desc->m_output_index).get_target_inputs()) {
+ input.replace_source_output(result->get_input_source_output(0));
+ }
+ } else {
+ // "Incorrect type of the output description."
+ return false;
+ }
+ }
+
+ return true;
+ };
+
+ auto m = std::make_shared<ngraph::pattern::Matcher>(tensor_iterator, "UnrollTensorIterator");
+ register_matcher(m, callback);
+}
\ No newline at end of file
--- /dev/null
+// Copyright (C) 2020 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include <gtest/gtest.h>
+
+#include "common_test_utils/test_common.hpp"
+#include <string>
+#include <memory>
+#include <queue>
+
+#include <ngraph/function.hpp>
+#include <ngraph/opsets/opset4.hpp>
+#include <ngraph/pass/manager.hpp>
+#include <ngraph_ops/fully_connected.hpp>
+#include <transformations/tensor_iterator_transformations/unroll_tensor_iterator.hpp>
+#include <transformations/utils/utils.hpp>
+#include <transformations/init_node_info.hpp>
+
+#include "common_test_utils/ngraph_test_utils.hpp"
+
+using namespace testing;
+using namespace ngraph;
+
+TEST(TransformationTests, UnrollTensorIteratorGRUCell) {
+ std::shared_ptr<ngraph::Function> f(nullptr), f_ref(nullptr);
+ {
+ auto X = std::make_shared<opset4::Parameter>(element::f32, Shape{2, 1, 16});
+ auto Y = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ auto Xi = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 1, 16});
+ auto Yi = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ // Body
+ auto axis = ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0});
+ auto squeeze = std::make_shared<opset4::Squeeze>(Xi, axis);
+
+ auto w_val = std::vector<float>(384*16, 0);
+ auto r_val = std::vector<float>(384*128, 0);
+ auto b_val = std::vector<float>(384, 0);
+ auto W = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{384, 16}, w_val);
+ auto R = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{384, 128}, r_val);
+ auto B = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{384}, b_val);
+
+ auto gru_cell = std::make_shared<opset4::GRUCell>(squeeze, Yi, W, R, B, 128);
+ auto res_1 = std::make_shared<opset4::Result>(gru_cell);
+ auto unsqueeze = std::make_shared<opset4::Unsqueeze>(gru_cell, axis);
+ auto res_2 = std::make_shared<opset4::Result>(unsqueeze);
+ auto body = std::make_shared<opset4::TensorIterator::BodyLambda>(OutputVector{res_1, res_2},
+ ParameterVector{Xi, Yi});
+
+ auto tensor_iterator = std::make_shared<opset4::TensorIterator>();
+ tensor_iterator->set_body(body);
+
+ tensor_iterator->set_sliced_input(Xi, X, 0, 1, 1, -1, 0);
+ tensor_iterator->set_merged_input(Yi, Y, res_1);
+
+ auto out0 = tensor_iterator->get_iter_value(res_1, -1);
+ auto out1 = tensor_iterator->get_concatenated_slices(res_2, 0, 1, 1, -1, 0);
+
+ auto res_ti_1 = std::make_shared<opset4::Result>(tensor_iterator->output(1));
+ //auto res_ti_2 = std::make_shared<opset4::Result>(tensor_iterator->output(0));
+ f = std::make_shared<ngraph::Function>(ngraph::NodeVector{res_ti_1},
+ ngraph::ParameterVector{X, Y});
+
+ ngraph::pass::Manager manager;
+ manager.register_pass<ngraph::pass::InitNodeInfo>();
+ manager.register_pass<ngraph::pass::UnrollTensorIterator>();
+ manager.run_passes(f);
+
+ ASSERT_NO_THROW(check_rt_info(f));
+ }
+
+ {
+ auto X = std::make_shared<opset4::Parameter>(element::f32, Shape{2, 1, 16});
+ auto Y = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ auto axis_split = ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0});
+ auto split = std::make_shared<opset4::Split>(X, axis_split, 2);
+ auto axis = ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0});
+ auto squeeze_1 = std::make_shared<opset4::Squeeze>(split->output(0), axis);
+ auto squeeze_2 = std::make_shared<opset4::Squeeze>(split->output(1), axis);
+
+ auto w_val = std::vector<float>(384*16, 0);
+ auto r_val = std::vector<float>(384*128, 0);
+ auto b_val = std::vector<float>(384, 0);
+ auto W = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{384, 16}, w_val);
+ auto R = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{384, 128}, r_val);
+ auto B = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{384}, b_val);
+
+ auto gru_cell_1 = std::make_shared<opset4::GRUCell>(squeeze_1, Y, W, R, B, 128);
+ auto gru_cell_2 = std::make_shared<opset4::GRUCell>(squeeze_2, gru_cell_1, W, R, B, 128);
+
+ auto unsqueeze_1 = std::make_shared<opset4::Unsqueeze>(gru_cell_1, axis);
+ auto unsqueeze_2 = std::make_shared<opset4::Unsqueeze>(gru_cell_2, axis);
+ auto concat = std::make_shared<opset4::Concat>(OutputVector{unsqueeze_1, unsqueeze_2}, 0);
+
+ auto res_ti_1 = std::make_shared<opset4::Result>(concat);
+ //auto res_ti_2 = std::make_shared<opset4::Result>(unsqueeze_2);
+ f_ref = std::make_shared<ngraph::Function>(ngraph::NodeVector{res_ti_1}, ngraph::ParameterVector{X, Y});
+ }
+
+ auto res = compare_functions(f, f_ref);
+ ASSERT_TRUE(res.first) << res.second;
+}
+
+TEST(TransformationTests, UnrollTensorIteratorRNNCell) {
+ std::shared_ptr<ngraph::Function> f(nullptr), f_ref(nullptr);
+ {
+ auto X = std::make_shared<opset4::Parameter>(element::f32, Shape{2, 1, 16});
+ auto Y = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ auto Xi = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 1, 16});
+ auto Yi = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ // Body
+ auto axis = ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0});
+ auto squeeze = std::make_shared<opset4::Squeeze>(Xi, axis);
+
+ auto w_val = std::vector<float>(128*16, 0);
+ auto r_val = std::vector<float>(128*128, 0);
+ auto b_val = std::vector<float>(128, 0);
+ auto W = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{128, 16}, w_val);
+ auto R = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{128, 128}, r_val);
+ auto B = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{128}, b_val);
+
+ auto rnn_cell = std::make_shared<opset4::RNNCell>(squeeze, Yi, W, R, B, 128);
+ auto res_1 = std::make_shared<opset4::Result>(rnn_cell);
+ auto unsqueeze = std::make_shared<opset4::Unsqueeze>(rnn_cell, axis);
+ auto res_2 = std::make_shared<opset4::Result>(unsqueeze);
+ auto body = std::make_shared<opset4::TensorIterator::BodyLambda>(OutputVector{res_1, res_2},
+ ParameterVector{Xi, Yi});
+
+ auto tensor_iterator = std::make_shared<opset4::TensorIterator>();
+ tensor_iterator->set_body(body);
+
+ tensor_iterator->set_sliced_input(Xi, X, 0, 1, 1, -1, 0);
+ tensor_iterator->set_merged_input(Yi, Y, res_1);
+
+ auto out0 = tensor_iterator->get_iter_value(res_1, -1);
+ auto out1 = tensor_iterator->get_concatenated_slices(res_2, 0, 1, 1, -1, 0);
+
+ auto res_ti_1 = std::make_shared<opset4::Result>(tensor_iterator->output(1));
+ //auto res_ti_2 = std::make_shared<opset4::Result>(tensor_iterator->output(0));
+ f = std::make_shared<ngraph::Function>(ngraph::NodeVector{res_ti_1},
+ ngraph::ParameterVector{X, Y});
+
+ ngraph::pass::Manager manager;
+ manager.register_pass<ngraph::pass::InitNodeInfo>();
+ manager.register_pass<ngraph::pass::UnrollTensorIterator>();
+ manager.run_passes(f);
+
+ ASSERT_NO_THROW(check_rt_info(f));
+ }
+
+ {
+ auto X = std::make_shared<opset4::Parameter>(element::f32, Shape{2, 1, 16});
+ auto Y = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ auto axis_split = ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0});
+ auto split = std::make_shared<opset4::Split>(X, axis_split, 2);
+ auto axis = ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0});
+ auto squeeze_1 = std::make_shared<opset4::Squeeze>(split->output(0), axis);
+ auto squeeze_2 = std::make_shared<opset4::Squeeze>(split->output(1), axis);
+
+ auto w_val = std::vector<float>(128*16, 0);
+ auto r_val = std::vector<float>(128*128, 0);
+ auto b_val = std::vector<float>(128, 0);
+ auto W = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{128, 16}, w_val);
+ auto R = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{128, 128}, r_val);
+ auto B = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{128}, b_val);
+
+ auto rnn_cell_1 = std::make_shared<opset4::RNNCell>(squeeze_1, Y, W, R, B, 128);
+ auto rnn_cell_2 = std::make_shared<opset4::RNNCell>(squeeze_2, rnn_cell_1, W, R, B, 128);
+
+ auto unsqueeze_1 = std::make_shared<opset4::Unsqueeze>(rnn_cell_1, axis);
+ auto unsqueeze_2 = std::make_shared<opset4::Unsqueeze>(rnn_cell_2, axis);
+ auto concat = std::make_shared<opset4::Concat>(OutputVector{unsqueeze_1, unsqueeze_2}, 0);
+
+ auto res_ti_1 = std::make_shared<opset4::Result>(concat);
+ //auto res_ti_2 = std::make_shared<opset4::Result>(unsqueeze_2);
+ f_ref = std::make_shared<ngraph::Function>(ngraph::NodeVector{res_ti_1}, ngraph::ParameterVector{X, Y});
+ }
+
+ auto res = compare_functions(f, f_ref);
+ ASSERT_TRUE(res.first) << res.second;
+}
+
+TEST(TransformationTests, UnrollTensorIteratorLSTMCell) {
+ std::shared_ptr<ngraph::Function> f(nullptr), f_ref(nullptr);
+ {
+ auto X = std::make_shared<opset4::Parameter>(element::f32, Shape{2, 1, 16});
+ auto Y = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+ auto Z = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ auto Xi = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 1, 16});
+ auto Yi = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+ auto Zi = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ // Body
+ auto axis = ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0});
+ auto squeeze = std::make_shared<opset4::Squeeze>(Xi, axis);
+
+ auto w_val = std::vector<float>(512*16, 0);
+ auto r_val = std::vector<float>(512*128, 0);
+ auto b_val = std::vector<float>(512, 0);
+ auto W = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{512, 16}, w_val);
+ auto R = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{512, 128}, r_val);
+ auto B = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{512}, b_val);
+
+ auto lstm_cell = std::make_shared<opset4::LSTMCell>(squeeze, Yi, Zi, W, R, B, 128);
+ auto res_1 = std::make_shared<opset4::Result>(lstm_cell);
+ auto unsqueeze = std::make_shared<opset4::Unsqueeze>(lstm_cell, axis);
+ auto res_2 = std::make_shared<opset4::Result>(unsqueeze);
+ auto body = std::make_shared<opset4::TensorIterator::BodyLambda>(OutputVector{res_1, res_2},
+ ParameterVector{Xi, Yi, Zi});
+
+ auto tensor_iterator = std::make_shared<opset4::TensorIterator>();
+ tensor_iterator->set_body(body);
+
+ tensor_iterator->set_invariant_input(Zi, Z);
+ tensor_iterator->set_sliced_input(Xi, X, 0, 1, 1, -1, 0);
+ tensor_iterator->set_merged_input(Yi, Y, res_1);
+
+ auto out0 = tensor_iterator->get_iter_value(res_1, -1);
+ auto out1 = tensor_iterator->get_concatenated_slices(res_2, 0, 1, 1, -1, 0);
+
+ auto res_ti_1 = std::make_shared<opset4::Result>(tensor_iterator->output(1));
+ //auto res_ti_2 = std::make_shared<opset4::Result>(tensor_iterator->output(0));
+ f = std::make_shared<ngraph::Function>(ngraph::NodeVector{res_ti_1},
+ ngraph::ParameterVector{X, Y, Z});
+
+ ngraph::pass::Manager manager;
+ manager.register_pass<ngraph::pass::InitNodeInfo>();
+ manager.register_pass<ngraph::pass::UnrollTensorIterator>();
+ manager.run_passes(f);
+
+ ASSERT_NO_THROW(check_rt_info(f));
+ }
+
+ {
+ auto X = std::make_shared<opset4::Parameter>(element::f32, Shape{2, 1, 16});
+ auto Y = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+ auto Z = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ auto axis_split = ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0});
+ auto split = std::make_shared<opset4::Split>(X, axis_split, 2);
+ auto axis = ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0});
+ auto squeeze_1 = std::make_shared<opset4::Squeeze>(split->output(0), axis);
+ auto squeeze_2 = std::make_shared<opset4::Squeeze>(split->output(1), axis);
+
+ auto w_val = std::vector<float>(512*16, 0);
+ auto r_val = std::vector<float>(512*128, 0);
+ auto b_val = std::vector<float>(512, 0);
+ auto W = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{512, 16}, w_val);
+ auto R = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{512, 128}, r_val);
+ auto B = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{512}, b_val);
+
+ auto lstm_cell_1 = std::make_shared<opset4::LSTMCell>(squeeze_1, Y, Z, W, R, B, 128);
+ auto lstm_cell_2 = std::make_shared<opset4::LSTMCell>(squeeze_2, lstm_cell_1, Z, W, R, B, 128);
+
+ auto unsqueeze_1 = std::make_shared<opset4::Unsqueeze>(lstm_cell_1, axis);
+ auto unsqueeze_2 = std::make_shared<opset4::Unsqueeze>(lstm_cell_2, axis);
+ auto concat = std::make_shared<opset4::Concat>(OutputVector{unsqueeze_1, unsqueeze_2}, 0);
+
+ auto res_ti_1 = std::make_shared<opset4::Result>(concat);
+ //auto res_ti_2 = std::make_shared<opset4::Result>(unsqueeze_2);
+ f_ref = std::make_shared<ngraph::Function>(ngraph::NodeVector{res_ti_1}, ngraph::ParameterVector{X, Y, Z});
+ }
+
+ auto res = compare_functions(f, f_ref);
+ ASSERT_TRUE(res.first) << res.second;
+}
+
+TEST(TransformationTests, UnrollTensorIteratorGRUCellSingleIteration) {
+ std::shared_ptr<ngraph::Function> f(nullptr), f_ref(nullptr);
+ {
+ auto X = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 1, 16});
+ auto Y = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ auto Xi = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 1, 16});
+ auto Yi = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ // Body
+ auto axis = ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0});
+ auto squeeze = std::make_shared<opset4::Squeeze>(Xi, axis);
+
+ auto w_val = std::vector<float>(384*16, 0);
+ auto r_val = std::vector<float>(384*128, 0);
+ auto b_val = std::vector<float>(384, 0);
+ auto W = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{384, 16}, w_val);
+ auto R = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{384, 128}, r_val);
+ auto B = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{384}, b_val);
+
+ auto gru_cell = std::make_shared<opset4::GRUCell>(squeeze, Yi, W, R, B, 128);
+ auto res_1 = std::make_shared<opset4::Result>(gru_cell);
+ auto unsqueeze = std::make_shared<opset4::Unsqueeze>(gru_cell, axis);
+ auto res_2 = std::make_shared<opset4::Result>(unsqueeze);
+ auto body = std::make_shared<opset4::TensorIterator::BodyLambda>(OutputVector{res_1, res_2},
+ ParameterVector{Xi, Yi});
+
+ auto tensor_iterator = std::make_shared<opset4::TensorIterator>();
+ tensor_iterator->set_body(body);
+
+ tensor_iterator->set_sliced_input(Xi, X, 0, 1, 1, -1, 0);
+ tensor_iterator->set_merged_input(Yi, Y, res_1);
+
+ auto out0 = tensor_iterator->get_iter_value(res_1, -1);
+ auto out1 = tensor_iterator->get_concatenated_slices(res_2, 0, 1, 1, -1, 0);
+
+ auto res_ti_1 = std::make_shared<opset4::Result>(tensor_iterator->output(1));
+ //auto res_ti_2 = std::make_shared<opset4::Result>(tensor_iterator->output(0));
+ f = std::make_shared<ngraph::Function>(ngraph::NodeVector{res_ti_1},
+ ngraph::ParameterVector{X, Y});
+
+ ngraph::pass::Manager manager;
+ manager.register_pass<ngraph::pass::InitNodeInfo>();
+ manager.register_pass<ngraph::pass::UnrollTensorIterator>();
+ manager.run_passes(f);
+
+ ASSERT_NO_THROW(check_rt_info(f));
+ }
+
+ {
+ auto X = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 1, 16});
+ auto Y = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ auto axis = ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0});
+ auto squeeze_1 = std::make_shared<opset4::Squeeze>(X, axis);
+
+ auto w_val = std::vector<float>(384*16, 0);
+ auto r_val = std::vector<float>(384*128, 0);
+ auto b_val = std::vector<float>(384, 0);
+ auto W = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{384, 16}, w_val);
+ auto R = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{384, 128}, r_val);
+ auto B = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{384}, b_val);
+
+ auto gru_cell_1 = std::make_shared<opset4::GRUCell>(squeeze_1, Y, W, R, B, 128);
+
+ auto unsqueeze_1 = std::make_shared<opset4::Unsqueeze>(gru_cell_1, axis);
+
+ auto res_ti_1 = std::make_shared<opset4::Result>(unsqueeze_1);
+ //auto res_ti_2 = std::make_shared<opset4::Result>(unsqueeze_2);
+ f_ref = std::make_shared<ngraph::Function>(ngraph::NodeVector{res_ti_1}, ngraph::ParameterVector{X, Y});
+ }
+
+ auto res = compare_functions(f, f_ref);
+ ASSERT_TRUE(res.first) << res.second;
+}
+
+TEST(TransformationTests, UnrollTensorIteratorRNNCellSingleIteration) {
+ std::shared_ptr<ngraph::Function> f(nullptr), f_ref(nullptr);
+ {
+ auto X = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 1, 16});
+ auto Y = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ auto Xi = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 1, 16});
+ auto Yi = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ // Body
+ auto axis = ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0});
+ auto squeeze = std::make_shared<opset4::Squeeze>(Xi, axis);
+
+ auto w_val = std::vector<float>(128*16, 0);
+ auto r_val = std::vector<float>(128*128, 0);
+ auto b_val = std::vector<float>(128, 0);
+ auto W = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{128, 16}, w_val);
+ auto R = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{128, 128}, r_val);
+ auto B = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{128}, b_val);
+
+ auto rnn_cell = std::make_shared<opset4::RNNCell>(squeeze, Yi, W, R, B, 128);
+ auto res_1 = std::make_shared<opset4::Result>(rnn_cell);
+ auto unsqueeze = std::make_shared<opset4::Unsqueeze>(rnn_cell, axis);
+ auto res_2 = std::make_shared<opset4::Result>(unsqueeze);
+ auto body = std::make_shared<opset4::TensorIterator::BodyLambda>(OutputVector{res_1, res_2},
+ ParameterVector{Xi, Yi});
+
+ auto tensor_iterator = std::make_shared<opset4::TensorIterator>();
+ tensor_iterator->set_body(body);
+
+ tensor_iterator->set_sliced_input(Xi, X, 0, 1, 1, -1, 0);
+ tensor_iterator->set_merged_input(Yi, Y, res_1);
+
+ auto out0 = tensor_iterator->get_iter_value(res_1, -1);
+ auto out1 = tensor_iterator->get_concatenated_slices(res_2, 0, 1, 1, -1, 0);
+
+ auto res_ti_1 = std::make_shared<opset4::Result>(tensor_iterator->output(1));
+ //auto res_ti_2 = std::make_shared<opset4::Result>(tensor_iterator->output(0));
+ f = std::make_shared<ngraph::Function>(ngraph::NodeVector{res_ti_1},
+ ngraph::ParameterVector{X, Y});
+
+ ngraph::pass::Manager manager;
+ manager.register_pass<ngraph::pass::InitNodeInfo>();
+ manager.register_pass<ngraph::pass::UnrollTensorIterator>();
+ manager.run_passes(f);
+
+ ASSERT_NO_THROW(check_rt_info(f));
+ }
+
+ {
+ auto X = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 1, 16});
+ auto Y = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ auto axis = ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0});
+ auto squeeze_1 = std::make_shared<opset4::Squeeze>(X, axis);
+
+ auto w_val = std::vector<float>(128*16, 0);
+ auto r_val = std::vector<float>(128*128, 0);
+ auto b_val = std::vector<float>(128, 0);
+ auto W = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{128, 16}, w_val);
+ auto R = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{128, 128}, r_val);
+ auto B = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{128}, b_val);
+
+ auto rnn_cell_1 = std::make_shared<opset4::RNNCell>(squeeze_1, Y, W, R, B, 128);
+
+ auto unsqueeze_1 = std::make_shared<opset4::Unsqueeze>(rnn_cell_1, axis);
+ auto res_ti_1 = std::make_shared<opset4::Result>(unsqueeze_1);
+
+ f_ref = std::make_shared<ngraph::Function>(ngraph::NodeVector{res_ti_1}, ngraph::ParameterVector{X, Y});
+ }
+
+ auto res = compare_functions(f, f_ref);
+ ASSERT_TRUE(res.first) << res.second;
+}
+
+TEST(TransformationTests, UnrollTensorIteratorLSTMCellSingleIterationSingleIteration) {
+ std::shared_ptr<ngraph::Function> f(nullptr), f_ref(nullptr);
+ {
+ auto X = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 1, 16});
+ auto Y = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+ auto Z = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ auto Xi = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 1, 16});
+ auto Yi = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+ auto Zi = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ // Body
+ auto axis = ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0});
+ auto squeeze = std::make_shared<opset4::Squeeze>(Xi, axis);
+
+ auto w_val = std::vector<float>(512*16, 0);
+ auto r_val = std::vector<float>(512*128, 0);
+ auto b_val = std::vector<float>(512, 0);
+ auto W = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{512, 16}, w_val);
+ auto R = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{512, 128}, r_val);
+ auto B = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{512}, b_val);
+
+ auto lstm_cell = std::make_shared<opset4::LSTMCell>(squeeze, Yi, Zi, W, R, B, 128);
+ auto res_1 = std::make_shared<opset4::Result>(lstm_cell);
+ auto unsqueeze = std::make_shared<opset4::Unsqueeze>(lstm_cell, axis);
+ auto res_2 = std::make_shared<opset4::Result>(unsqueeze);
+ auto body = std::make_shared<opset4::TensorIterator::BodyLambda>(OutputVector{res_1, res_2},
+ ParameterVector{Xi, Yi, Zi});
+
+ auto tensor_iterator = std::make_shared<opset4::TensorIterator>();
+ tensor_iterator->set_body(body);
+
+ tensor_iterator->set_invariant_input(Zi, Z);
+ tensor_iterator->set_sliced_input(Xi, X, 0, 1, 1, -1, 0);
+ tensor_iterator->set_merged_input(Yi, Y, res_1);
+
+ auto out0 = tensor_iterator->get_iter_value(res_1, -1);
+ auto out1 = tensor_iterator->get_concatenated_slices(res_2, 0, 1, 1, -1, 0);
+
+ auto res_ti_1 = std::make_shared<opset4::Result>(tensor_iterator->output(1));
+ //auto res_ti_2 = std::make_shared<opset4::Result>(tensor_iterator->output(0));
+ f = std::make_shared<ngraph::Function>(ngraph::NodeVector{res_ti_1},
+ ngraph::ParameterVector{X, Y, Z});
+
+ ngraph::pass::Manager manager;
+ manager.register_pass<ngraph::pass::InitNodeInfo>();
+ manager.register_pass<ngraph::pass::UnrollTensorIterator>();
+ manager.run_passes(f);
+
+ ASSERT_NO_THROW(check_rt_info(f));
+ }
+
+ {
+ auto X = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 1, 16});
+ auto Y = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+ auto Z = std::make_shared<opset4::Parameter>(element::f32, Shape{1, 128});
+
+ auto axis = ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0});
+ auto squeeze_1 = std::make_shared<opset4::Squeeze>(X, axis);
+
+ auto w_val = std::vector<float>(512*16, 0);
+ auto r_val = std::vector<float>(512*128, 0);
+ auto b_val = std::vector<float>(512, 0);
+ auto W = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{512, 16}, w_val);
+ auto R = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{512, 128}, r_val);
+ auto B = ngraph::opset4::Constant::create(ngraph::element::f32, ngraph::Shape{512}, b_val);
+
+ auto lstm_cell_1 = std::make_shared<opset4::LSTMCell>(squeeze_1, Y, Z, W, R, B, 128);
+
+ auto unsqueeze_1 = std::make_shared<opset4::Unsqueeze>(lstm_cell_1, axis);
+ auto res_ti_1 = std::make_shared<opset4::Result>(unsqueeze_1);
+ //auto res_ti_2 = std::make_shared<opset4::Result>(unsqueeze_2);
+ f_ref = std::make_shared<ngraph::Function>(ngraph::NodeVector{res_ti_1}, ngraph::ParameterVector{X, Y, Z});
+ }
+
+ auto res = compare_functions(f, f_ref);
+ ASSERT_TRUE(res.first) << res.second;
+}
projects / platform / upstream / dldt.git / commitdiff