text: // Copyright (C) 2018-2021 Intel Corporation
  // SPDX-License-Identifier: Apache-2.0
  //
  
  #include <gtest/gtest.h>
  
  #include <string>
  #include <memory>
  
  #include <ngraph/function.hpp>
  #include <ngraph/opsets/opset2.hpp>
  #include <ngraph/opsets/opset3.hpp>
  #include <transformations/op_conversions/convert_shuffle_channels3.hpp>
  #include <transformations/init_node_info.hpp>
  #include <ngraph/pass/manager.hpp>
  
  #include "common_test_utils/ngraph_test_utils.hpp"
  
  using namespace testing;
  using namespace ngraph;
  
  std::shared_ptr<ngraph::Function> buildInputGraph(int64_t axis, int64_t group, const ::PartialShape& p) {
      auto input = std::make_shared<::opset3::Parameter>(::element::f32, p);
      auto shuffle_channels = std::make_shared<::opset3::ShuffleChannels>(input, axis, group);
      return std::make_shared<::Function>(::NodeVector{shuffle_channels}, ::ParameterVector{input});
  }
  
  TEST_F(TransformationTestsF, ConvertShuffleChannelsAxis0) {
      int64_t group = 4;
      auto ps = ::PartialShape{12, Dimension::dynamic(), Dimension::dynamic(), Dimension::dynamic()};
      function = buildInputGraph(0, group, ps);
      manager.register_pass<ngraph::pass::ConvertShuffleChannels3>();
  
      auto input = std::make_shared<::opset3::Parameter>(::element::f32, ps);
  
      auto reduce_axis_const = ::opset2::Constant::create(element::i64, Shape({1}), std::vector<int64_t>{0});
      auto original_shape = std::make_shared<::opset2::ShapeOf>(input->output(0));
      auto split_input_dimensions = std::make_shared<::opset2::VariadicSplit>(
              original_shape->output(0),
              ::opset2::Constant::create(element::i64, Shape({1}), std::vector<int64_t>{0}),
              ::opset2::Constant::create(element::i64, Shape({2}), {1, 3}));
  
      ::OutputVector new_dims = {
              ::opset2::Constant::create(element::i64, Shape({1}), std::vector<int64_t>{group}),
              ::opset2::Constant::create(element::i64, Shape({1}), std::vector<int64_t>{-1}),
              std::make_shared<::opset2::ReduceProd>(split_input_dimensions->output(1), reduce_axis_const, true)};
  
      auto new_shape = std::make_shared<::opset2::Concat>(new_dims, 0);
      auto reshape = std::make_shared<::opset2::Reshape>(input->output(0), new_shape, false);
      auto transpose = std::make_shared<::opset2::Transpose>(reshape->output(0),
                                                             ::opset2::Constant::create(element::i64, Shape({3}),
                                                                                        {1, 0, 2}));
      auto reshape_back = std::make_shared<::opset2::Reshape>(transpose->output(0), original_shape->output(0), false);
  
      function_ref = std::make_shared<::Function>(::NodeVector{reshape_back}, ::ParameterVector{input});
  }
  
  TEST_F(TransformationTestsF, ConvertShuffleChannelsAxis1) {
      int64_t group = 4;
      auto ps = ::PartialShape{Dimension::dynamic(), 12, Dimension::dynamic(), Dimension::dynamic()};
      function = buildInputGraph(1, group, ps);
      manager.register_pass<ngraph::pass::ConvertShuffleChannels3>();
  
      auto input = std::make_shared<::opset3::Parameter>(::element::f32, ps);
  
      auto reduce_axis_const = ::opset2::Constant::create(element::i64, Shape({1}), std::vector<int64_t>{0});
      auto original_shape = std::make_shared<::opset2::ShapeOf>(input->output(0));
      auto split_input_dimensions = std::make_shared<::opset2::VariadicSplit>(
              original_shape->output(0),
              ::opset2::Constant::create(element::i64, Shape({1}), std::vector<int64_t>{0}),
              ::opset2::Constant::create(element::i64, Shape({3}), {1, 1, 2}));
  
      ::OutputVector new_dims = {
              std::make_shared<::opset2::ReduceProd>(split_input_dimensions->output(0), reduce_axis_const, true),
              ::opset2::Constant::create(element::i64, Shape({1}), std::vector<int64_t>{group}),
              ::opset2::Constant::create(element::i64, Shape({1}), std::vector<int64_t>{-1}),
              std::make_shared<::opset2::ReduceProd>(split_input_dimensions->output(2), reduce_axis_const, true)};
  
      auto new_shape = std::make_shared<::opset2::Concat>(new_dims, 0);
      auto reshape = std::make_shared<::opset2::Reshape>(input->output(0), new_shape, false);
      auto transpose = std::make_shared<::opset2::Transpose>(reshape->output(0),
                                                             ::opset2::Constant::create(element::i64, Shape({4}),
                                                                                        {0, 2, 1, 3}));
      auto reshape_back = std::make_shared<::opset2::Reshape>(transpose->output(0), original_shape->output(0), false);
  
      function_ref = std::make_shared<::Function>(::NodeVector{reshape_back}, ::ParameterVector{input});
  }
  
  TEST_F(TransformationTestsF, ConvertShuffleChannelsAxis2) {
      int64_t group = 4;
      auto ps = ::PartialShape{Dimension::dynamic(), Dimension::dynamic(), 12, Dimension::dynamic()};
      function = buildInputGraph(2, group, ps);
      manager.register_pass<ngraph::pass::ConvertShuffleChannels3>();
  
      auto input = std::make_shared<::opset3::Parameter>(::element::f32, ps);
  
      auto reduce_axis_const = ::opset2::Constant::create(element::i64, Shape({1}), std::vector<int64_t>{0});
      auto original_shape = std::make_shared<::opset2::ShapeOf>(input->output(0));
      auto split_input_dimensions = std::make_shared<::opset2::VariadicSplit>(
              original_shape->output(0),
              ::opset2::Constant::create(element::i64, Shape({1}), std::vector<int64_t>{0}),
              ::opset2::Constant::create(element::i64, Shape({3}), {2, 1, 1}));
  
      ::OutputVector new_dims = {
              std::make_shared<::opset2::ReduceProd>(split_input_dimensions->output(0), reduce_axis_const, true),
              ::opset2::Constant::create(element::i64, Shape({1}), std::vector<int64_t>{group}),
              ::opset2::Constant::create(element::i64, Shape({1}), std::vector<int64_t>{-1}),
              std::make_shared<::opset2::ReduceProd>(split_input_dimensions->output(2), reduce_axis_const, true)};
  
      auto new_shape = std::make_shared<::opset2::Concat>(new_dims, 0);
      auto reshape = std::make_shared<::opset2::Reshape>(input->output(0), new_shape, false);
      auto transpose = std::make_shared<::opset2::Transpose>(reshape->output(0),
                                                             ::opset2::Constant::create(element::i64, Shape({4}),
                                                                                        {0, 2, 1, 3}));
      auto reshape_back = std::make_shared<::opset2::Reshape>(transpose->output(0), original_shape->output(0), false);
  
      function_ref = std::make_shared<::Function>(::NodeVector{reshape_back}, ::ParameterVector{input});
  }
  
  TEST_F(TransformationTestsF, ConvertShuffleChannelsLastAxis) {
      int64_t group = 4;
      auto ps = ::PartialShape{Dimension::dynamic(), Dimension::dynamic(), Dimension::dynamic(), 12};
      function = buildInputGraph(-1, group, ps);
      manager.register_pass<ngraph::pass::ConvertShuffleChannels3>();
  
      auto input = std::make_shared<::opset3::Parameter>(::element::f32, ps);
  
      auto reduce_axis_const = ::opset2::Constant::create(element::i64, Shape({1}), std::vector<int64_t>{0});
      auto original_shape = std::make_shared<::opset2::ShapeOf>(input->output(0));
      auto split_input_dimensions = std::make_shared<::opset2::VariadicSplit>(
              original_shape->output(0),
              ::opset2::Constant::create(element::i64, Shape({1}), std::vector<int64_t>{0}),
              ::opset2::Constant::create(element::i64, Shape({2}), {3, 1}));
  
      ::OutputVector new_dims = {
              std::make_shared<::opset2::ReduceProd>(split_input_dimensions->output(0), reduce_axis_const, true),
              ::opset2::Constant::create(element::i64, Shape({1}), std::vector<int64_t>{group}),
              ::opset2::Constant::create(element::i64, Shape({1}), std::vector<int64_t>{-1})};
  
      auto new_shape = std::make_shared<::opset2::Concat>(new_dims, 0);
      auto reshape = std::make_shared<::opset2::Reshape>(input->output(0), new_shape, false);
      auto transpose = std::make_shared<::opset2::Transpose>(reshape->output(0),
                                                             ::opset2::Constant::create(element::i64, Shape({3}),
                                                                                        {0, 2, 1}));
      auto reshape_back = std::make_shared<::opset2::Reshape>(transpose->output(0), original_shape->output(0), false);
  
      function_ref = std::make_shared<::Function>(::NodeVector{reshape_back}, ::ParameterVector{input});
  }