text: //*****************************************************************************
  // Copyright 2017-2020 Intel Corporation
  //
  // Licensed under the Apache License, Version 2.0 (the "License");
  // you may not use this file except in compliance with the License.
  // You may obtain a copy of the License at
  //
  //     http://www.apache.org/licenses/LICENSE-2.0
  //
  // Unless required by applicable law or agreed to in writing, software
  // distributed under the License is distributed on an "AS IS" BASIS,
  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  // See the License for the specific language governing permissions and
  // limitations under the License.
  //*****************************************************************************
  #include <numeric>
  
  #include "ngraph/builder/split.hpp"
  #include "ngraph/op/constant.hpp"
  #include "ngraph/op/fused/split.hpp"
  #include "ngraph/validation_util.hpp"
  
  using namespace std;
  using namespace ngraph;
  
  constexpr NodeTypeInfo op::v0::Split::type_info;
  
  op::v0::Split::Split(const Output<Node>& data, const Output<Node>& axis, const size_t num_split)
      : FusedOp({data, axis})
      , m_split_evenly{true}
      , m_num_split{num_split}
  {
      constructor_validate_and_infer_types();
  }
  
  op::v0::Split::Split(const Output<Node>& data,
                       const Output<Node>& axis,
                       const std::vector<size_t>& splits)
      : FusedOp({data, axis})
      , m_split_evenly{false}
      , m_num_split{0}
      , m_splits{splits}
  {
      constructor_validate_and_infer_types();
  }
  
  void op::v0::Split::pre_validate_and_infer_types()
  {
      const auto axis_shape = input(1).get_shape();
      NODE_VALIDATION_CHECK(this, is_scalar(axis_shape), "The 'axis' input node must be scalar");
  
      const auto axis_node = input_value(1).get_node_shared_ptr();
      NODE_VALIDATION_CHECK(this, axis_node->is_constant(), "The 'axis' input node must be constant");
      const auto axis_node_const = as_type_ptr<op::Constant>(axis_node);
      m_axis = axis_node_const->get_vector<int64_t>()[0];
  
      // Create dynamic-typed outputs. Actual shape/type will be computed during shape inference
      for (size_t i = 0; i < std::max(m_splits.size(), m_num_split); i++)
      {
          set_output_type(i, input(0).get_element_type(), PartialShape::dynamic());
      }
  
      if (is_dynamic())
      {
          return;
      }
  
      const auto shape = input(0).get_shape();
  
      m_axis = ngraph::normalize_axis(this, m_axis, shape.size());
  
      const auto dimension_at_axis = shape.at(m_axis);
      if (m_split_evenly)
      {
          NODE_VALIDATION_CHECK(this,
                                dimension_at_axis % m_num_split == 0,
                                "The input tensor's dimension pointed by the 'axis' parameter: ",
                                dimension_at_axis,
                                " has to be a multiple of the 'num_split' parameter value: ",
                                m_num_split);
  
          m_splits.assign(m_num_split, dimension_at_axis / m_num_split);
      }
      else
      {
          const auto sum_splits = accumulate(begin(m_splits), end(m_splits), 0UL);
          NODE_VALIDATION_CHECK(this,
                                sum_splits == dimension_at_axis,
                                "The input tensor's dimension pointed by the 'axis' parameter: ",
                                dimension_at_axis,
                                " has to be equal to the sum of splits passed to the op: ",
                                sum_splits);
  
          const bool all_splits_positive =
              all_of(begin(m_splits), end(m_splits), [](const size_t v) { return v > 0; });
  
          NODE_VALIDATION_CHECK(this,
                                all_splits_positive == true,
                                "All values of the 'splits' attribute must be greater than zero");
      }
      set_input_is_relevant_to_shape(0);
  }
  
  NodeVector op::v0::Split::decompose_op() const
  {
      return builder::split(input_value(0), m_splits, m_axis);
  }
  
  shared_ptr<Node> op::v0::Split::copy_with_new_args(const NodeVector& new_args) const
  {
      check_new_args_count(this, new_args);
      return make_shared<Split>(new_args.at(0), new_args.at(1), m_splits);
  }
  
  constexpr NodeTypeInfo op::v1::Split::type_info;
  
  op::v1::Split::Split(const Output<Node>& data, const Output<Node>& axis, const size_t num_splits)
      : FusedOp({data, axis})
      , m_num_splits{num_splits}
  {
      constructor_validate_and_infer_types();
  }
  
  void op::v1::Split::validate_and_infer_types()
  {
      const auto data_ps = input_value(0).get_partial_shape();
      const auto axis_ps = input_value(1).get_partial_shape();
      const auto axis_et = input_value(1).get_element_type();
  
      NODE_VALIDATION_CHECK(this,
                            axis_ps.rank().is_static() && (size_t)axis_ps.rank() == 0,
                            "The 'axis' input is expected to be a scalar. Got: ",
                            axis_ps);
  
      NODE_VALIDATION_CHECK(
          this, axis_et.is_integral(), "The 'axis' input only accepts integral types");
  
      if (input_value(1).get_node_shared_ptr()->is_constant())
      {
          const auto axis_input = as_type_ptr<op::Constant>(input_value(1).get_node_shared_ptr());
          auto axis = axis_input->cast_vector<int64_t>()[0];
  
          if (data_ps.is_static())
          {
              const auto data_shape = data_ps.to_shape();
              axis = ngraph::normalize_axis(this, axis, data_shape.size());
  
              const auto dimension_at_axis = data_shape.at(axis);
  
              NODE_VALIDATION_CHECK(this,
                                    dimension_at_axis % m_num_splits == 0,
                                    "The input tensor's dimension pointed by the 'axis' parameter: ",
                                    dimension_at_axis,
                                    " has to be a multiple of the 'num_splits' attribute value: ",
                                    m_num_splits);
  
              Shape each_output_shape{data_shape};
              each_output_shape.at(axis) = dimension_at_axis / m_num_splits;
  
              for (size_t i = 0; i < m_num_splits; ++i)
              {
                  set_output_type(i, input(0).get_element_type(), each_output_shape);
              }
          }
      }
      else
      {
          for (size_t i = 0; i < m_num_splits; ++i)
          {
              set_output_type(i, input(0).get_element_type(), PartialShape::dynamic());
          }
  
          set_input_is_relevant_to_shape(0);
      }
  }
  
  shared_ptr<Node> op::v1::Split::copy_with_new_args(const NodeVector& new_args) const
  {
      check_new_args_count(this, new_args);
      return make_shared<v1::Split>(new_args.at(0), new_args.at(1), m_num_splits);
  }
  