ngraph/core/include/ngraph/op/split.hpp

   1 //*****************************************************************************
   2 // Copyright 2017-2020 Intel Corporation
   3 //
   4 // Licensed under the Apache License, Version 2.0 (the "License");
   5 // you may not use this file except in compliance with the License.
   6 // You may obtain a copy of the License at
   7 //
   8 //     http://www.apache.org/licenses/LICENSE-2.0
   9 //
  10 // Unless required by applicable law or agreed to in writing, software
  11 // distributed under the License is distributed on an "AS IS" BASIS,
  12 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 // See the License for the specific language governing permissions and
  14 // limitations under the License.
  15 //*****************************************************************************
  16
  17 #pragma once
  18
  19 #include <memory>
  20 #include <vector>
  21
  22 #include "ngraph/node.hpp"
  23 #include "ngraph/op/util/fused_op.hpp"
  24
  25 NGRAPH_SUPPRESS_DEPRECATED_START
  26
  27 namespace ngraph
  28 {
  29     namespace op
  30     {
  31         namespace v0
  32         {
  33             /// \brief Splits the input tensor into a list of smaller tensors ("pieces")
  34             class NGRAPH_DEPRECATED(
  35                 "This operation is deprecated and will be removed soon. "
  36                 "Use v1::Split instead of it.") NGRAPH_API Split : public ngraph::op::util::FusedOp
  37             {
  38                 NGRAPH_SUPPRESS_DEPRECATED_START
  39             public:
  40                 static constexpr NodeTypeInfo type_info{"Split", 0};
  41                 const NodeTypeInfo& get_type_info() const override { return type_info; }
  42                 Split() = default;
  43                 /// \brief Constructs a Split op that evenly divides the input tensor.
  44                 ///
  45                 /// \param data       Node producing the input tensor
  46                 /// \param axis       Node producing an axis along which the input tensor
  47                 ///                   should be split. Negative values mean counting from
  48                 ///                   the back of the input tensor's shape.
  49                 /// \param num_split  a number of "pieces" the input tensor will be split to
  50                 Split(const Output<Node>& data, const Output<Node>& axis, const size_t num_split);
  51
  52                 /// \brief Constructs a Split op that splits the input tensor into variable length
  53                 ///        "pieces"
  54                 ///
  55                 /// \param data    Node producing the input tensor
  56                 /// \param axis    Node producing an axis along which the input tensor
  57                 ///                should be split. Negative values mean counting from
  58                 ///                the back of the input tensor's shape.
  59                 /// \param splits  a list of lengths that the input tensor should be
  60                 ///                split to. Use this constructor to split the input
  61                 ///                tensor to variable length chunks.
  62                 Split(const Output<Node>& data,
  63                       const Output<Node>& axis,
  64                       const std::vector<size_t>& splits);
  65
  66                 void pre_validate_and_infer_types() override;
  67
  68                 virtual OutputVector decompose_op() const override;
  69
  70                 virtual std::shared_ptr<Node>
  71                     clone_with_new_inputs(const OutputVector& new_args) const override;
  72
  73                 size_t get_axis() const { return m_axis; }
  74                 const std::vector<size_t>& get_splits() const { return m_splits; }
  75             private:
  76                 /// used internally for validation purposes, indicates which constructor was used
  77                 bool m_split_evenly;
  78                 int64_t m_axis;
  79                 size_t m_num_split;
  80                 /// contains lengths of chunks that the input tensor will be split into
  81                 std::vector<size_t> m_splits;
  82                 NGRAPH_SUPPRESS_DEPRECATED_END
  83             };
  84         }
  85
  86         namespace v1
  87         {
  88             /// \brief Splits the input tensor into a list of equal sized tensors
  89             class NGRAPH_API Split : public ngraph::op::Op
  90             {
  91             public:
  92                 static constexpr NodeTypeInfo type_info{"Split", 1};
  93                 const NodeTypeInfo& get_type_info() const override { return type_info; }
  94                 /// \brief Constructs a split operation.
  95                 Split() = default;
  96                 /// \brief Constructs a split operation.
  97                 /// \param data        The tensor to be split.
  98                 /// \param axis        The index of an axis in "data" along which to perform
  99                 ///                    the split.
 100                 /// \param num_splits  The number of pieces that the data tensor should be
 101                 ///                    split into.
 102                 Split(const Output<Node>& data, const Output<Node>& axis, const size_t num_splits);
 103
 104                 bool visit_attributes(AttributeVisitor& visitor) override;
 105                 void validate_and_infer_types() override;
 106                 virtual std::shared_ptr<Node>
 107                     clone_with_new_inputs(const OutputVector& new_args) const override;
 108
 109                 size_t get_num_splits() const { return m_num_splits; }
 110                 void set_num_splits(const size_t num_splits) { m_num_splits = num_splits; }
 111                 bool evaluate(const HostTensorVector& outputs,
 112                               const HostTensorVector& inputs) const override;
 113
 114             protected:
 115                 size_t m_num_splits;
 116             };
 117         }
 118
 119         NGRAPH_SUPPRESS_DEPRECATED_START
 120         using v0::Split;
 121         NGRAPH_SUPPRESS_DEPRECATED_END
 122     }
 123 }
 124
 125 NGRAPH_SUPPRESS_DEPRECATED_END