Skip to content
This repository was archived by the owner on Jan 3, 2023. It is now read-only.

Commit dd8c9ed

Browse files
ayzhuangdiyessi
ayzhuang
authored and
diyessi
committed
Enable Gather with negative indices. (#3701)
* Enable Gather with negative indices. * Address PR feedback. * Remove GOE from Gather unit tests.
1 parent b8266ca commit dd8c9ed

File tree

5 files changed

+70
-31
lines changed

5 files changed

+70
-31
lines changed

‎src/ngraph/runtime/cpu/kernel/gather.hpp

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,8 @@ namespace ngraph
6969
Eigen::array<Eigen::Index, Rank1> in_dims;
7070
Eigen::array<Eigen::Index, Rank2> out_dims;
7171

72+
auto axis_length = inputs_shape[axis];
73+
7274
for (int i = 0; i < Rank1; i++)
7375
{
7476
in_dims[i] = inputs_shape[i];
@@ -84,6 +86,7 @@ namespace ngraph
8486
static_cast<ElementType*>(inputs), in_dims);
8587

8688
auto indices_ptr = static_cast<IndicesType*>(indices);
89+
IndicesType index_value;
8790
auto indices_rank = indices_shape.size();
8891
auto outer_loop_num = 1;
8992
for (int i = 0; i < axis; i++)
@@ -123,7 +126,10 @@ namespace ngraph
123126
// at axis
124127
in_extents[axis] = 1;
125128
// at axis, get the value from indices arg
126-
in_offsets[axis] = indices_ptr[0];
129+
index_value = indices_ptr[0];
130+
// take care of negative indices
131+
in_offsets[axis] =
132+
index_value >= 0 ? index_value : index_value + axis_length;
127133

128134
// before axis
129135
for (int r = 0; r < axis; r++)
@@ -193,7 +199,10 @@ namespace ngraph
193199
}
194200
// at axis, get the value from indices arg
195201
int k = i % num_indices;
196-
in_offsets[axis] = indices_ptr[k];
202+
index_value = indices_ptr[k];
203+
// take care of negative indices
204+
in_offsets[axis] =
205+
index_value >= 0 ? index_value : index_value + axis_length;
197206

198207
// indices_from_indices_arg depends on indices_shape and k.
199208
// suppose the inputs has shape {3, 3, 3}, indices has shape {2, 2}, and

‎src/ngraph/runtime/intelgpu/unit_test.manifest

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ gather_2d_indices_axis_1_2d_input
6464
gather_scalar_indices_no_axis_2d_input
6565
gather_1d_indices_no_axis_1d_input
6666
gather_2d_indices_no_axis_2d_input
67+
gather_2d_negative_and_positive_indices_no_axis_2d_input
6768
gather_3d_indices_no_axis_2d_input
6869
gather_4d_indices_no_axis_2d_input
6970
gemm

‎src/ngraph/runtime/plaidml/unit_test.manifest

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,6 +196,7 @@ gather_4d_indices_no_axis_uint8
196196
gather_4d_indices_no_axis_2d_input
197197
gather_3d_indices_no_axis_2d_input
198198
gather_2d_indices_no_axis_2d_input
199+
gather_2d_negative_and_positive_indices_no_axis_2d_input
199200
gather_1d_indices_no_axis_1d_input
200201
gather_scalar_indices_no_axis_2d_input
201202
gather_2d_indices_axis_1_2d_input

‎src/ngraph/runtime/reference/gather_nd.hpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,8 @@ namespace ngraph
8383
for (size_t i = 0; i < slice_rank; i++)
8484
{
8585
U index = indices[indices_index];
86+
// take care of negative indices
87+
index = index >= 0 ? index : index + params_shape[i];
8688
params_start_corner[i] = index;
8789
params_end_corner[i] = index + 1;
8890
indices_index++;

‎test/backend/gather.in.cpp

Lines changed: 55 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_4d_indices_no_axis_uint8)
4444
auto P = make_shared<op::Parameter>(element::u8, params_shape);
4545
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
4646
auto G = make_shared<op::Gather>(P, I);
47-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
47+
auto f = make_shared<Function>(G, ParameterVector{P, I});
4848

4949
auto backend = runtime::Backend::create("${BACKEND_NAME}");
5050

@@ -77,7 +77,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_4d_indices_no_axis_2d_input)
7777
auto P = make_shared<op::Parameter>(element::f32, params_shape);
7878
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
7979
auto G = make_shared<op::Gather>(P, I);
80-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
80+
auto f = make_shared<Function>(G, ParameterVector{P, I});
8181

8282
auto backend = runtime::Backend::create("${BACKEND_NAME}");
8383

@@ -113,7 +113,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_3d_indices_no_axis_2d_input)
113113
auto P = make_shared<op::Parameter>(element::f32, params_shape);
114114
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
115115
auto G = make_shared<op::Gather>(P, I);
116-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
116+
auto f = make_shared<Function>(G, ParameterVector{P, I});
117117

118118
auto backend = runtime::Backend::create("${BACKEND_NAME}");
119119

@@ -144,7 +144,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_2d_indices_no_axis_2d_input)
144144
auto P = make_shared<op::Parameter>(element::f32, params_shape);
145145
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
146146
auto G = make_shared<op::Gather>(P, I);
147-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
147+
auto f = make_shared<Function>(G, ParameterVector{P, I});
148148

149149
auto backend = runtime::Backend::create("${BACKEND_NAME}");
150150

@@ -162,6 +162,32 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_2d_indices_no_axis_2d_input)
162162
MIN_FLOAT_TOLERANCE_BITS));
163163
}
164164

165+
NGRAPH_TEST(${BACKEND_NAME}, gather_2d_negative_and_positive_indices_no_axis_2d_input)
166+
{
167+
Shape params_shape{3, 2};
168+
Shape indices_shape{2, 2};
169+
Shape out_shape{2, 2, 2};
170+
auto P = make_shared<op::Parameter>(element::f32, params_shape);
171+
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
172+
auto G = make_shared<op::Gather>(P, I);
173+
auto f = make_shared<Function>(G, ParameterVector{P, I});
174+
175+
auto backend = runtime::Backend::create("${BACKEND_NAME}");
176+
177+
// Create some tensors for input/output
178+
auto p = backend->create_tensor(element::f32, params_shape);
179+
copy_data(p, vector<float>{1.0f, 1.1f, 2.0f, 2.1f, 3.0f, 3.1f});
180+
auto i = backend->create_tensor(element::i32, indices_shape);
181+
copy_data(i, vector<int32_t>{0, -2, 1, 2});
182+
auto result = backend->create_tensor(element::f32, out_shape);
183+
184+
auto c = backend->compile(f);
185+
c->call_with_validate({result}, {p, i});
186+
EXPECT_TRUE(test::all_close_f((vector<float>{1.0f, 1.1f, 2.0f, 2.1f, 2.0f, 2.1f, 3.0f, 3.1f}),
187+
read_vector<float>(result),
188+
MIN_FLOAT_TOLERANCE_BITS));
189+
}
190+
165191
NGRAPH_TEST(${BACKEND_NAME}, gather_1d_indices_no_axis_1d_input)
166192
{
167193
Shape params_shape{3};
@@ -170,7 +196,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_1d_indices_no_axis_1d_input)
170196
auto P = make_shared<op::Parameter>(element::f32, params_shape);
171197
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
172198
auto G = make_shared<op::Gather>(P, I);
173-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
199+
auto f = make_shared<Function>(G, ParameterVector{P, I});
174200

175201
auto backend = runtime::Backend::create("${BACKEND_NAME}");
176202

@@ -195,7 +221,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_scalar_indices_no_axis_2d_input)
195221
auto P = make_shared<op::Parameter>(element::f32, params_shape);
196222
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
197223
auto G = make_shared<op::Gather>(P, I);
198-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
224+
auto f = make_shared<Function>(G, ParameterVector{P, I});
199225

200226
auto backend = runtime::Backend::create("${BACKEND_NAME}");
201227

@@ -220,7 +246,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_2d_indices_axis_1_2d_input)
220246
auto P = make_shared<op::Parameter>(element::f32, params_shape);
221247
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
222248
auto G = make_shared<op::Gather>(P, I, 1);
223-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
249+
auto f = make_shared<Function>(G, ParameterVector{P, I});
224250

225251
auto backend = runtime::Backend::create("${BACKEND_NAME}");
226252

@@ -246,7 +272,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_1d_indices_axis_2_4d_input)
246272
auto P = make_shared<op::Parameter>(element::f32, params_shape);
247273
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
248274
auto G = make_shared<op::Gather>(P, I, 2);
249-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
275+
auto f = make_shared<Function>(G, ParameterVector{P, I});
250276

251277
auto backend = runtime::Backend::create("${BACKEND_NAME}");
252278

@@ -277,7 +303,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_scalar_indices_axis_1_2d_input)
277303
auto P = make_shared<op::Parameter>(element::f32, params_shape);
278304
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
279305
auto G = make_shared<op::Gather>(P, I, 1);
280-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
306+
auto f = make_shared<Function>(G, ParameterVector{P, I});
281307

282308
auto backend = runtime::Backend::create("${BACKEND_NAME}");
283309

@@ -302,7 +328,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_single_indices)
302328
auto P = make_shared<op::Parameter>(element::f32, params_shape);
303329
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
304330
auto G = make_shared<op::GatherND>(P, I);
305-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
331+
auto f = make_shared<Function>(G, ParameterVector{P, I});
306332

307333
auto backend = runtime::Backend::create("${BACKEND_NAME}");
308334

@@ -327,7 +353,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_scalar_from_2d)
327353
auto P = make_shared<op::Parameter>(element::f32, params_shape);
328354
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
329355
auto G = make_shared<op::GatherND>(P, I);
330-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
356+
auto f = make_shared<Function>(G, ParameterVector{P, I});
331357

332358
auto backend = runtime::Backend::create("${BACKEND_NAME}");
333359

@@ -352,7 +378,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_1d_from_2d)
352378
auto P = make_shared<op::Parameter>(element::f32, params_shape);
353379
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
354380
auto G = make_shared<op::GatherND>(P, I);
355-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
381+
auto f = make_shared<Function>(G, ParameterVector{P, I});
356382

357383
auto backend = runtime::Backend::create("${BACKEND_NAME}");
358384

@@ -378,7 +404,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_scalar_from_3d)
378404
auto P = make_shared<op::Parameter>(element::f32, params_shape);
379405
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
380406
auto G = make_shared<op::GatherND>(P, I);
381-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
407+
auto f = make_shared<Function>(G, ParameterVector{P, I});
382408

383409
auto backend = runtime::Backend::create("${BACKEND_NAME}");
384410

@@ -403,7 +429,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_1d_from_3d)
403429
auto P = make_shared<op::Parameter>(element::f32, params_shape);
404430
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
405431
auto G = make_shared<op::GatherND>(P, I);
406-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
432+
auto f = make_shared<Function>(G, ParameterVector{P, I});
407433

408434
auto backend = runtime::Backend::create("${BACKEND_NAME}");
409435

@@ -429,7 +455,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_2d_from_3d)
429455
auto P = make_shared<op::Parameter>(element::f32, params_shape);
430456
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
431457
auto G = make_shared<op::GatherND>(P, I);
432-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
458+
auto f = make_shared<Function>(G, ParameterVector{P, I});
433459

434460
auto backend = runtime::Backend::create("${BACKEND_NAME}");
435461

@@ -455,7 +481,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_batch_scalar_from_2d)
455481
auto P = make_shared<op::Parameter>(element::f32, params_shape);
456482
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
457483
auto G = make_shared<op::GatherND>(P, I);
458-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
484+
auto f = make_shared<Function>(G, ParameterVector{P, I});
459485

460486
auto backend = runtime::Backend::create("${BACKEND_NAME}");
461487

@@ -480,7 +506,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_batch_1d_from_2d)
480506
auto P = make_shared<op::Parameter>(element::f32, params_shape);
481507
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
482508
auto G = make_shared<op::GatherND>(P, I);
483-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
509+
auto f = make_shared<Function>(G, ParameterVector{P, I});
484510

485511
auto backend = runtime::Backend::create("${BACKEND_NAME}");
486512

@@ -506,7 +532,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_batch_scalar_from_3d)
506532
auto P = make_shared<op::Parameter>(element::f32, params_shape);
507533
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
508534
auto G = make_shared<op::GatherND>(P, I);
509-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
535+
auto f = make_shared<Function>(G, ParameterVector{P, I});
510536

511537
auto backend = runtime::Backend::create("${BACKEND_NAME}");
512538

@@ -532,7 +558,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_batch_1d_from_3d)
532558
auto P = make_shared<op::Parameter>(element::f32, params_shape);
533559
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
534560
auto G = make_shared<op::GatherND>(P, I);
535-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
561+
auto f = make_shared<Function>(G, ParameterVector{P, I});
536562

537563
auto backend = runtime::Backend::create("${BACKEND_NAME}");
538564

@@ -558,7 +584,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_batch_2d_from_3d)
558584
auto P = make_shared<op::Parameter>(element::f32, params_shape);
559585
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
560586
auto G = make_shared<op::GatherND>(P, I);
561-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
587+
auto f = make_shared<Function>(G, ParameterVector{P, I});
562588

563589
auto backend = runtime::Backend::create("${BACKEND_NAME}");
564590

@@ -584,7 +610,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_int8)
584610
auto P = make_shared<op::Parameter>(element::i8, params_shape);
585611
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
586612
auto G = make_shared<op::Gather>(P, I);
587-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
613+
auto f = make_shared<Function>(G, ParameterVector{P, I});
588614

589615
auto backend = runtime::Backend::create("${BACKEND_NAME}");
590616

@@ -610,7 +636,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_int16)
610636
auto P = make_shared<op::Parameter>(element::i16, params_shape);
611637
auto I = make_shared<op::Parameter>(element::i64, indices_shape);
612638
auto G = make_shared<op::Gather>(P, I);
613-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
639+
auto f = make_shared<Function>(G, ParameterVector{P, I});
614640

615641
auto backend = runtime::Backend::create("${BACKEND_NAME}");
616642

@@ -636,7 +662,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_int32)
636662
auto P = make_shared<op::Parameter>(element::i32, params_shape);
637663
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
638664
auto G = make_shared<op::Gather>(P, I);
639-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
665+
auto f = make_shared<Function>(G, ParameterVector{P, I});
640666

641667
auto backend = runtime::Backend::create("${BACKEND_NAME}");
642668

@@ -662,7 +688,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_int64)
662688
auto P = make_shared<op::Parameter>(element::i64, params_shape);
663689
auto I = make_shared<op::Parameter>(element::i64, indices_shape);
664690
auto G = make_shared<op::Gather>(P, I);
665-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
691+
auto f = make_shared<Function>(G, ParameterVector{P, I});
666692

667693
auto backend = runtime::Backend::create("${BACKEND_NAME}");
668694

@@ -688,7 +714,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_uint8)
688714
auto P = make_shared<op::Parameter>(element::u8, params_shape);
689715
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
690716
auto G = make_shared<op::Gather>(P, I);
691-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
717+
auto f = make_shared<Function>(G, ParameterVector{P, I});
692718

693719
auto backend = runtime::Backend::create("${BACKEND_NAME}");
694720

@@ -714,7 +740,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_uint16)
714740
auto P = make_shared<op::Parameter>(element::u16, params_shape);
715741
auto I = make_shared<op::Parameter>(element::i64, indices_shape);
716742
auto G = make_shared<op::Gather>(P, I);
717-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
743+
auto f = make_shared<Function>(G, ParameterVector{P, I});
718744

719745
auto backend = runtime::Backend::create("${BACKEND_NAME}");
720746

@@ -740,7 +766,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_uint32)
740766
auto P = make_shared<op::Parameter>(element::u32, params_shape);
741767
auto I = make_shared<op::Parameter>(element::i32, indices_shape);
742768
auto G = make_shared<op::Gather>(P, I);
743-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
769+
auto f = make_shared<Function>(G, ParameterVector{P, I});
744770

745771
auto backend = runtime::Backend::create("${BACKEND_NAME}");
746772

@@ -766,7 +792,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_uint64)
766792
auto P = make_shared<op::Parameter>(element::u64, params_shape);
767793
auto I = make_shared<op::Parameter>(element::i64, indices_shape);
768794
auto G = make_shared<op::Gather>(P, I);
769-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
795+
auto f = make_shared<Function>(G, ParameterVector{P, I});
770796

771797
auto backend = runtime::Backend::create("${BACKEND_NAME}");
772798

@@ -792,7 +818,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_bool)
792818
auto P = make_shared<op::Parameter>(element::boolean, params_shape);
793819
auto I = make_shared<op::Parameter>(element::i64, indices_shape);
794820
auto G = make_shared<op::Gather>(P, I);
795-
auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
821+
auto f = make_shared<Function>(G, ParameterVector{P, I});
796822

797823
auto backend = runtime::Backend::create("${BACKEND_NAME}");
798824

0 commit comments

Comments
 (0)