1 //*****************************************************************************
2 // Copyright 2017-2020 Intel Corporation
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
8 // http://www.apache.org/licenses/LICENSE-2.0
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 //*****************************************************************************
21 #include "ngraph/axis_set.hpp"
22 #include "ngraph/op/constant.hpp"
23 #include "ngraph/op/op.hpp"
31 // \brief Computes indices of top k maximum/minimum index along a specified axis for a
33 class NGRAPH_DEPRECATED(
34 "This operation is deprecated and will be removed soon. "
35 "Use v1::TopK instead of it.") NGRAPH_API TopK : public Op
37 NGRAPH_SUPPRESS_DEPRECATED_START
39 using SortType = TopKSortType;
41 static constexpr NodeTypeInfo type_info{"TopK", 0};
42 const NodeTypeInfo& get_type_info() const override { return type_info; }
43 /// \brief Constructs a TopK operation
45 /// \brief Constructs a TopK operation.
47 /// \param arg The input tensor
48 /// \param top_k_axis The axis along which to compute top k indices
49 /// \param index_element_type produce indices. Currently, only int64 or int32 are
51 /// \param k Number of top indices to compute. Compute all indices if k = 0
52 /// \param compute_max Compute top k max or top k min?
53 /// \param sort SortType for sorting results, default - SORT_VALUES
54 TopK(const Output<Node>& arg,
56 const element::Type& index_element_type,
58 bool compute_max = true,
59 SortType sort = SortType::SORT_VALUES);
60 /// \brief Constructs a TopK operation.
62 /// \param arg The input tensor
63 /// \param k Number of top indices to compute. Compute all indices if k = 0
64 /// \param top_k_axis The axis along which to compute top k indices
65 /// \param index_element_type produce indices. Currently, only int64 or int32 are
67 /// \param compute_max Compute top k max or top k min?
68 /// \param sort SortType for sorting results, default - SORT_VALUES
69 TopK(const Output<Node>& arg,
70 const Output<Node>& k,
72 const element::Type& index_element_type,
73 bool compute_max = true,
74 SortType sort = SortType::SORT_VALUES);
76 /// \brief Constructs a TopK operation.
78 /// \param arg The input tensor
79 /// \param k Number of top indices to compute. Compute all indices if k = 0
80 /// \param top_k_axis The axis along which to compute top k indices
81 /// \param index_element_type produce indices. Currently, only int64 or int32 are
83 /// \param compute_max Compute top k max or top k min?
84 /// \param sort SortType for sorting results, default - NONE
85 TopK(const Output<Node>& arg,
86 const Output<Node>& k,
87 const Output<Node>& top_k_axis,
88 const element::Type& index_element_type,
89 bool compute_max = true,
90 SortType sort = SortType::NONE);
92 void validate_and_infer_types() override;
94 virtual std::shared_ptr<Node>
95 clone_with_new_inputs(const OutputVector& new_args) const override;
100 size_t get_top_k_axis() const;
101 Dimension get_top_k_axis_dynamic() const;
102 void set_top_k_axis(size_t k);
104 element::Type get_index_element_type() const { return m_index_element_type; }
105 bool get_compute_max() const { return m_compute_max; }
106 SortType get_sort() const { return m_sort; }
107 size_t get_default_output_index() const override { return no_default_index(); }
108 bool evaluate(const HostTensorVector& outputs,
109 const HostTensorVector& inputs) const override;
112 element::Type m_index_element_type;
113 bool m_compute_max{false};
114 SortType m_sort{SortType::NONE};
115 Shape compute_output_shape(const Shape input_shape,
117 const size_t axis) const;
118 NGRAPH_SUPPRESS_DEPRECATED_END
124 /// \brief Computes indices and values of the k maximum/minimum values
125 /// for each slice along specified axis.
126 class NGRAPH_API TopK : public Op
129 using SortType = TopKSortType;
130 using Mode = TopKMode;
132 static constexpr NodeTypeInfo type_info{"TopK", 1};
133 const NodeTypeInfo& get_type_info() const override { return type_info; }
134 /// \brief Constructs a TopK operation
136 /// \brief Constructs a TopK operation with two outputs: values and indices.
137 /// By default the indices output is described by i32 data type.
139 /// \param data The input tensor
140 /// \param k Specifies how many maximum/minimum elements should be computed
141 /// (note: scalar input tensor)
142 /// \param axis The axis along which to compute top k indices
143 /// \param mode Specifies which operation (min or max) is used to select
144 /// the biggest element of two.
145 /// \param sort Specifies order of output elements and/or indices
146 /// Accepted values: none, index, value
147 /// \param index_element_type Specyfies type of produced indices
148 TopK(const Output<Node>& data,
149 const Output<Node>& k,
151 const std::string& mode,
152 const std::string& sort,
153 const element::Type& index_element_type = element::i32);
155 TopK(const Output<Node>& data,
156 const Output<Node>& k,
160 const element::Type& index_element_type = element::i32);
162 bool visit_attributes(AttributeVisitor& visitor) override;
163 void validate_and_infer_types() override;
165 virtual std::shared_ptr<Node>
166 clone_with_new_inputs(const OutputVector& new_args) const override;
168 virtual size_t get_version() const override { return 1; }
169 /// \brief Returns axis value after normalization
170 /// \note If input rank required to normalization is dynamic, the exception is
172 uint64_t get_axis() const;
173 /// \brief Returns axis value before normalization
174 int64_t get_provided_axis() const { return m_axis; }
175 void set_axis(const int64_t axis);
176 Mode get_mode() const { return m_mode; }
177 void set_mode(const Mode mode) { m_mode = mode; }
178 SortType get_sort_type() const { return m_sort; }
179 void set_sort_type(const SortType sort) { m_sort = sort; }
180 element::Type get_index_element_type() const { return m_index_element_type; }
181 void set_index_element_type(const element::Type& index_element_type)
183 m_index_element_type = index_element_type;
185 /// \brief Returns the value of K, if available
187 /// \note If the second input to this op is a constant, the value is retrieved
188 /// and returned. If the input is not constant(dynamic) this method returns 0
189 size_t get_k() const;
190 void set_k(size_t k);
191 size_t get_default_output_index() const override { return no_default_index(); }
192 bool evaluate(const HostTensorVector& outputs,
193 const HostTensorVector& inputs) const override;
197 uint64_t m_normalized_axis;
200 element::Type m_index_element_type{element::i32};
202 virtual size_t read_k_from_constant_node(const std::shared_ptr<Node>& node,
203 const element::Type& k_element_type) const;
205 template <typename T>
206 size_t validate_and_get_k(const std::shared_ptr<op::Constant>& k_constant) const;
207 Shape compute_output_shape(const std::string& node_description,
208 const PartialShape input_partial_shape,
209 const int64_t k) const;
210 void set_axis(const Rank input_rank, const int64_t axis);
216 /// \brief Computes indices and values of the k maximum/minimum values
217 /// for each slice along specified axis.
218 class NGRAPH_API TopK : public v1::TopK
221 static constexpr NodeTypeInfo type_info{"TopK", 3};
222 const NodeTypeInfo& get_type_info() const override { return type_info; }
223 /// \brief Constructs a TopK operation
225 /// \brief Constructs a TopK operation with two outputs: values and indices.
226 /// By default the indices output is described by i32 data type.
228 /// \param data The input tensor
229 /// \param k Specifies how many maximum/minimum elements should be computed
230 /// (note: scalar input tensor)
231 /// \param axis The axis along which to compute top k indices
232 /// \param mode Specifies which operation (min or max) is used to select
233 /// the biggest element of two.
234 /// \param sort Specifies order of output elements and/or indices
235 /// Accepted values: none, index, value
236 /// \param index_element_type Specyfies type of produced indices
237 TopK(const Output<Node>& data,
238 const Output<Node>& k,
240 const std::string& mode,
241 const std::string& sort,
242 const element::Type& index_element_type = element::i32);
244 TopK(const Output<Node>& data,
245 const Output<Node>& k,
249 const element::Type& index_element_type = element::i32);
250 bool visit_attributes(AttributeVisitor& visitor) override;
251 void validate_and_infer_types() override;
252 virtual std::shared_ptr<Node>
253 clone_with_new_inputs(const OutputVector& new_args) const override;
255 bool evaluate(const HostTensorVector& outputs,
256 const HostTensorVector& inputs) const override;
260 read_k_from_constant_node(const std::shared_ptr<Node>& node,
261 const element::Type& k_element_type) const override;
265 NGRAPH_SUPPRESS_DEPRECATED_START
267 NGRAPH_SUPPRESS_DEPRECATED_END