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 //*****************************************************************************
17 #include "ngraph/op/gather.hpp"
19 #include "ngraph/op/constant.hpp"
20 #include "ngraph/runtime/host_tensor.hpp"
21 #include "ngraph/runtime/reference/gather.hpp"
22 #include "ngraph/shape.hpp"
26 NGRAPH_SUPPRESS_DEPRECATED_START
29 using namespace ngraph;
31 static const int PARAMS = 0;
32 static const int INDICES = 1;
33 static const int AXIS = 2;
35 constexpr NodeTypeInfo op::v0::Gather::type_info;
37 op::v0::Gather::Gather(const Output<Node>& params, const Output<Node>& indices, size_t axis)
38 : Op({params, indices})
41 constructor_validate_and_infer_types();
44 shared_ptr<Node> op::v0::Gather::clone_with_new_inputs(const OutputVector& new_args) const
46 check_new_args_count(this, new_args);
47 return make_shared<v0::Gather>(new_args.at(PARAMS), new_args.at(INDICES), m_axis);
50 void op::v0::Gather::validate_and_infer_types()
52 element::Type result_et = get_input_element_type(PARAMS);
53 element::Type indices_et = get_input_element_type(INDICES);
55 const PartialShape& params_shape = get_input_partial_shape(PARAMS);
56 const PartialShape& indices_shape = get_input_partial_shape(INDICES);
58 NODE_VALIDATION_CHECK(this,
59 indices_et == element::i32 || indices_et == element::i64,
60 "Indices element type must be i64 or i32");
62 // params rank must be at least (axis + 1)
63 // indices value must be in range [0, params.shape[axis]).
64 // output rank is rank(params) + rank(indices) - 1
65 NODE_VALIDATION_CHECK(this,
66 params_shape.rank().is_dynamic() ||
67 params_shape.rank().get_length() > static_cast<size_t>(m_axis),
68 "params rank is expected to be at least axis + 1");
70 PartialShape result_shape;
71 if (params_shape.rank().is_static() && indices_shape.rank().is_static())
73 std::vector<Dimension> result_dims(params_shape.rank().get_length() +
74 indices_shape.rank().get_length() - 1);
76 for (; i < static_cast<size_t>(m_axis); i++)
78 result_dims[i] = params_shape[i];
80 for (size_t j = 0; j < indices_shape.rank().get_length(); i++, j++)
82 result_dims[i] = indices_shape[j];
84 for (size_t j = static_cast<size_t>(m_axis) + 1; j < params_shape.rank().get_length();
87 result_dims[i] = params_shape[j];
90 result_shape = PartialShape(result_dims);
94 result_shape = PartialShape::dynamic();
97 set_output_type(0, result_et, result_shape);
100 constexpr NodeTypeInfo op::v1::Gather::type_info;
101 const int64_t op::v1::Gather::AXIS_NOT_SET_VALUE;
103 op::v1::Gather::Gather(const Output<Node>& params,
104 const Output<Node>& indices,
105 const Output<Node>& axes)
106 : Op({params, indices, axes})
108 constructor_validate_and_infer_types();
111 bool ngraph::op::v1::Gather::visit_attributes(AttributeVisitor& visitor)
116 void op::v1::Gather::validate_and_infer_types()
118 const auto& input_rank = get_input_partial_shape(PARAMS).rank();
119 const auto& axis_shape = get_input_partial_shape(AXIS);
120 const auto& axis_rank = axis_shape.rank();
122 if (axis_rank.is_static() && axis_shape.is_static())
124 const auto axis_is_scalar = axis_rank.get_length() == 0;
125 const auto axis_has_one_elem =
126 axis_rank.get_length() == 1 && axis_shape[0].get_length() == 1;
127 NODE_VALIDATION_CHECK(this,
128 axis_is_scalar || axis_has_one_elem,
129 "Axes input must be scalar or have 1 element (shape: ",
134 int64_t axis = get_axis();
135 if (input_rank.is_static() && axis != AXIS_NOT_SET_VALUE)
137 NODE_VALIDATION_CHECK(this,
138 axis < input_rank.get_length(),
139 "The axis must => 0 and <= input_rank (axis: ",
144 element::Type result_et = get_input_element_type(PARAMS);
145 element::Type indices_et = get_input_element_type(INDICES);
147 const PartialShape& params_shape = get_input_partial_shape(PARAMS);
148 const PartialShape& indices_shape = get_input_partial_shape(INDICES);
150 PartialShape result_shape;
151 if (params_shape.rank().is_static() && indices_shape.rank().is_static() &&
152 axis != AXIS_NOT_SET_VALUE)
154 std::vector<Dimension> result_dims(params_shape.rank().get_length() +
155 indices_shape.rank().get_length() - 1);
157 for (; i < axis; i++)
159 result_dims[i] = params_shape[i];
161 for (uint64_t j = 0; j < indices_shape.rank().get_length(); i++, j++)
163 result_dims[i] = indices_shape[j];
165 for (uint64_t j = axis + 1; j < params_shape.rank().get_length(); i++, j++)
167 result_dims[i] = params_shape[j];
170 result_shape = PartialShape(result_dims);
174 result_shape = PartialShape::dynamic();
177 set_output_type(0, result_et, result_shape);
180 int64_t op::v1::Gather::get_axis() const
182 int64_t axis = AXIS_NOT_SET_VALUE;
183 auto axes_input_node = input_value(AXIS).get_node_shared_ptr();
184 if (auto const_op = as_type_ptr<op::Constant>(axes_input_node))
186 axis = const_op->cast_vector<int64_t>()[0];
190 const auto& input_rank = get_input_partial_shape(PARAMS).rank();
191 if (input_rank.is_static())
193 axis += input_rank.get_length();
199 shared_ptr<Node> op::v1::Gather::clone_with_new_inputs(const OutputVector& new_args) const
201 check_new_args_count(this, new_args);
202 return make_shared<v1::Gather>(new_args.at(PARAMS), new_args.at(INDICES), new_args.at(AXIS));
207 template <element::Type_t ET>
208 bool evaluate(const HostTensorPtr& arg0,
209 const HostTensorPtr& arg1,
210 const HostTensorPtr& out,
213 using T = typename element_type_traits<ET>::value_type;
214 Shape params_shape = arg0->get_shape();
215 Shape indices_shape = arg1->get_shape();
216 Shape out_shape(params_shape.size() + indices_shape.size() - 1);
218 for (; i < axis; i++)
220 out_shape[i] = params_shape[i];
222 for (uint64_t j = 0; j < indices_shape.size(); i++, j++)
224 out_shape[i] = indices_shape[j];
226 for (uint64_t j = axis + 1; j < params_shape.size(); i++, j++)
228 out_shape[i] = params_shape[j];
231 out->set_shape(out_shape);
233 if (arg1->get_element_type() == element::i64)
235 runtime::reference::gather<T, int64_t>(arg0->get_data_ptr<ET>(),
236 arg1->get_data_ptr<int64_t>(),
237 out->get_data_ptr<ET>(),
243 else if (arg1->get_element_type() == element::i32)
245 runtime::reference::gather<T, int32_t>(arg0->get_data_ptr<ET>(),
246 arg1->get_data_ptr<int32_t>(),
247 out->get_data_ptr<ET>(),
255 throw ngraph_error("Unexpected type");
261 bool evaluate_gather(const HostTensorPtr& arg0,
262 const HostTensorPtr& arg1,
263 const HostTensorPtr& out,
268 switch (out->get_element_type())
270 TYPE_CASE(i32)(arg0, arg1, out, axis);
272 TYPE_CASE(i64)(arg0, arg1, out, axis);
274 TYPE_CASE(u32)(arg0, arg1, out, axis);
276 TYPE_CASE(u64)(arg0, arg1, out, axis);
278 TYPE_CASE(f16)(arg0, arg1, out, axis);
280 TYPE_CASE(f32)(arg0, arg1, out, axis);
282 TYPE_CASE(boolean)(arg0, arg1, out, axis);
284 default: rc = false; break;
290 bool op::v0::Gather::evaluate(const HostTensorVector& outputs, const HostTensorVector& inputs) const
292 OV_ITT_SCOPED_TASK(itt::domains::nGraphOp, "op::v0::Gather::evaluate");
293 return evaluate_gather(inputs[0], inputs[1], outputs[0], get_axis());
296 bool op::v1::Gather::evaluate(const HostTensorVector& outputs, const HostTensorVector& inputs) const
298 OV_ITT_SCOPED_TASK(itt::domains::nGraphOp, "op::v1::Gather::evaluate");
300 switch (inputs[2]->get_element_type())
302 case element::Type_t::i8: axis = inputs[2]->get_data_ptr<element::Type_t::i8>()[0]; break;
303 case element::Type_t::i16: axis = inputs[2]->get_data_ptr<element::Type_t::i16>()[0]; break;
304 case element::Type_t::i32: axis = inputs[2]->get_data_ptr<element::Type_t::i32>()[0]; break;
305 case element::Type_t::i64: axis = inputs[2]->get_data_ptr<element::Type_t::i64>()[0]; break;
306 case element::Type_t::u8: axis = inputs[2]->get_data_ptr<element::Type_t::u8>()[0]; break;
307 case element::Type_t::u16: axis = inputs[2]->get_data_ptr<element::Type_t::u16>()[0]; break;
308 case element::Type_t::u32: axis = inputs[2]->get_data_ptr<element::Type_t::u32>()[0]; break;
309 case element::Type_t::u64: axis = inputs[2]->get_data_ptr<element::Type_t::u64>()[0]; break;
310 default: throw ngraph_error("axis element type is not integral data type");
315 const auto& input_rank = get_input_partial_shape(PARAMS).rank();
316 if (input_rank.is_static())
318 axis += input_rank.get_length();
321 return evaluate_gather(inputs[0], inputs[1], outputs[0], axis);