ngraph/test/backend/topk.in.cpp

   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 #include <algorithm>
  18 #include <cinttypes>
  19 #include <cmath>
  20 #include <cstdlib>
  21 #include <numeric>
  22 #include <random>
  23 #include <string>
  24
  25 #include "gtest/gtest.h"
  26 #include "ngraph/op/constant.hpp"
  27 #include "ngraph/op/parameter.hpp"
  28 #include "ngraph/op/result.hpp"
  29 #include "ngraph/op/topk.hpp"
  30 #include "ngraph/runtime/tensor.hpp"
  31 #include "runtime/backend.hpp"
  32 #include "util/all_close_f.hpp"
  33 #include "util/test_control.hpp"
  34 #include "util/test_tools.hpp"
  35
  36 NGRAPH_SUPPRESS_DEPRECATED_START
  37
  38 using namespace std;
  39 using namespace ngraph;
  40
  41 static string s_manifest = "${MANIFEST}";
  42
  43 template <typename T>
  44 bool compare_set(const vector<T>& a, vector<T> b)
  45 {
  46     for (auto ita = a.begin(); ita != a.end(); ++ita)
  47     {
  48         auto itb = find(b.begin(), b.end(), *ita);
  49         if (itb == b.end())
  50         {
  51             return false;
  52         }
  53         else
  54         {
  55             b.erase(itb);
  56         }
  57     }
  58     return true;
  59 }
  60
  61 NGRAPH_TEST(${BACKEND_NAME}, topk_resnet50)
  62 {
  63     Shape shape{128, 1000};
  64     Shape rshape5{128, 5};
  65     Shape rshape1{128, 1};
  66     auto A = make_shared<op::Parameter>(element::f32, shape);
  67     auto B = make_shared<op::TopK>(A, 1, element::i32, 5, true);
  68     auto C = make_shared<op::TopK>(A, 1, element::i32, 1, true);
  69     auto out5_value = B->output(1);
  70     auto out5_index = B->output(0);
  71     auto out1_value = C->output(1);
  72     auto out1_index = C->output(0);
  73     auto f = make_shared<Function>(OutputVector{out5_value, out5_index, out1_value, out1_index},
  74                                    ParameterVector{A});
  75
  76     auto backend = runtime::Backend::create("${BACKEND_NAME}");
  77
  78     // Create some tensors for input/output
  79     auto a = backend->create_tensor(element::f32, shape);
  80     vector<float> data;
  81     for (size_t i = 0; i < shape[0]; i++)
  82     {
  83         for (size_t j = 0; j < shape[1]; j++)
  84         {
  85             data.push_back(j);
  86         }
  87     }
  88     copy_data(a, data);
  89
  90     auto result5_value = backend->create_tensor(element::f32, rshape5);
  91     auto result5_index = backend->create_tensor(element::i32, rshape5);
  92     auto result1_value = backend->create_tensor(element::f32, rshape1);
  93     auto result1_index = backend->create_tensor(element::i32, rshape1);
  94
  95     auto exec = backend->compile(f);
  96     exec->call({result5_value, result5_index, result1_value, result1_index}, {a});
  97
  98     auto actual5_value = read_vector<float>(result5_value);
  99     auto actual5_index = read_vector<int32_t>(result5_index);
 100     auto actual1_value = read_vector<float>(result1_value);
 101     auto actual1_index = read_vector<int32_t>(result1_index);
 102
 103     vector<float> expected5_value;
 104     vector<int32_t> expected5_index;
 105     for (size_t i = 0; i < rshape5[0]; i++)
 106     {
 107         for (size_t j = 0; j < rshape5[1]; j++)
 108         {
 109             expected5_value.push_back(shape[1] - j - 1);
 110             expected5_index.push_back(shape[1] - j - 1);
 111         }
 112     }
 113
 114     vector<float> expected1_value;
 115     vector<int32_t> expected1_index;
 116     for (size_t i = 0; i < rshape1[0]; i++)
 117     {
 118         for (size_t j = 0; j < rshape1[1]; j++)
 119         {
 120             expected1_value.push_back(shape[1] - j - 1);
 121             expected1_index.push_back(shape[1] - j - 1);
 122         }
 123     }
 124
 125     EXPECT_TRUE(compare_set<float>(expected5_value, actual5_value));
 126     EXPECT_TRUE(compare_set<int32_t>(expected5_index, actual5_index));
 127     EXPECT_TRUE(compare_set<float>(expected1_value, actual1_value));
 128     EXPECT_TRUE(compare_set<int32_t>(expected1_index, actual1_index));
 129 }
 130
 131 NGRAPH_TEST(${BACKEND_NAME}, topk_max_sort_none)
 132 {
 133     Shape shape{128, 1000};
 134     Shape rshape{128, 5};
 135     auto A = make_shared<op::Parameter>(element::f32, shape);
 136     auto B = make_shared<op::TopK>(A, 1, element::i32, 5, true, op::TopK::SortType::NONE);
 137     auto out_value = B->output(1);
 138     auto out_index = B->output(0);
 139     auto f = make_shared<Function>(OutputVector{out_value, out_index}, ParameterVector{A});
 140
 141     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 142
 143     // Create some tensors for input/output
 144     auto a = backend->create_tensor(element::f32, shape);
 145     vector<float> data;
 146     for (size_t i = 0; i < shape[0]; i++)
 147     {
 148         for (size_t j = 0; j < shape[1]; j++)
 149         {
 150             data.push_back(j);
 151         }
 152     }
 153     copy_data(a, data);
 154
 155     auto result_value = backend->create_tensor(element::f32, rshape);
 156     auto result_index = backend->create_tensor(element::i32, rshape);
 157
 158     auto exec = backend->compile(f);
 159     exec->call({result_value, result_index}, {a});
 160
 161     auto actual_value = read_vector<float>(result_value);
 162     auto actual_index = read_vector<int32_t>(result_index);
 163
 164     for (size_t i = 0; i < rshape[0]; i++)
 165     {
 166         vector<float> expected_value;
 167         vector<int32_t> expected_index;
 168         vector<float> act_value;
 169         vector<int32_t> act_index;
 170         for (size_t j = 0; j < rshape[1]; j++)
 171         {
 172             expected_value.push_back(shape[1] - j - 1);
 173             expected_index.push_back(shape[1] - j - 1);
 174             act_value.push_back(actual_value[rshape[1] * i + j]);
 175             act_index.push_back(actual_index[rshape[1] * i + j]);
 176         }
 177         EXPECT_TRUE(compare_set<float>(expected_value, act_value));
 178         EXPECT_TRUE(compare_set<int32_t>(expected_index, act_index));
 179     }
 180 }
 181
 182 NGRAPH_TEST(${BACKEND_NAME}, topk_min_sort_none)
 183 {
 184     Shape shape{128, 1000};
 185     Shape rshape{128, 5};
 186     auto A = make_shared<op::Parameter>(element::f32, shape);
 187     auto B = make_shared<op::TopK>(A, 1, element::i32, 5, false, op::TopK::SortType::NONE);
 188     auto out_value = B->output(1);
 189     auto out_index = B->output(0);
 190     auto f = make_shared<Function>(OutputVector{out_value, out_index}, ParameterVector{A});
 191
 192     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 193
 194     // Create some tensors for input/output
 195     auto a = backend->create_tensor(element::f32, shape);
 196     vector<float> data;
 197     for (size_t i = 0; i < shape[0]; i++)
 198     {
 199         for (size_t j = 0; j < shape[1]; j++)
 200         {
 201             data.push_back(j);
 202         }
 203     }
 204     copy_data(a, data);
 205
 206     auto result_value = backend->create_tensor(element::f32, rshape);
 207     auto result_index = backend->create_tensor(element::i32, rshape);
 208
 209     auto exec = backend->compile(f);
 210     exec->call({result_value, result_index}, {a});
 211
 212     auto actual_value = read_vector<float>(result_value);
 213     auto actual_index = read_vector<int32_t>(result_index);
 214
 215     for (size_t i = 0; i < rshape[0]; i++)
 216     {
 217         vector<float> expected_value;
 218         vector<int32_t> expected_index;
 219         vector<float> act_value;
 220         vector<int32_t> act_index;
 221         for (size_t j = 0; j < rshape[1]; j++)
 222         {
 223             expected_value.push_back(j);
 224             expected_index.push_back(j);
 225             act_value.push_back(actual_value[rshape[1] * i + j]);
 226             act_index.push_back(actual_index[rshape[1] * i + j]);
 227         }
 228         EXPECT_TRUE(compare_set<float>(expected_value, act_value));
 229         EXPECT_TRUE(compare_set<int32_t>(expected_index, act_index));
 230     }
 231 }
 232
 233 NGRAPH_TEST(${BACKEND_NAME}, topk_max_sort_value)
 234 {
 235     Shape shape{128, 1000};
 236     Shape rshape{128, 5};
 237     auto A = make_shared<op::Parameter>(element::f32, shape);
 238     auto B = make_shared<op::TopK>(A, 1, element::i32, 5, true, op::TopK::SortType::SORT_VALUES);
 239     auto out_value = B->output(1);
 240     auto out_index = B->output(0);
 241     auto f = make_shared<Function>(OutputVector{out_value, out_index}, ParameterVector{A});
 242
 243     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 244
 245     // Create some tensors for input/output
 246     auto a = backend->create_tensor(element::f32, shape);
 247     vector<float> data;
 248     for (size_t i = 0; i < shape[0]; i++)
 249     {
 250         for (size_t j = 0; j < shape[1]; j++)
 251         {
 252             data.push_back(j);
 253         }
 254     }
 255     copy_data(a, data);
 256
 257     auto result_value = backend->create_tensor(element::f32, rshape);
 258     auto result_index = backend->create_tensor(element::i32, rshape);
 259
 260     auto exec = backend->compile(f);
 261     exec->call({result_value, result_index}, {a});
 262
 263     auto actual_value = read_vector<float>(result_value);
 264     auto actual_index = read_vector<int32_t>(result_index);
 265
 266     vector<float> expected_value;
 267     vector<int32_t> expected_index;
 268     for (size_t i = 0; i < rshape[0]; i++)
 269     {
 270         for (size_t j = 0; j < rshape[1]; j++)
 271         {
 272             expected_value.push_back(shape[1] - j - 1);
 273             expected_index.push_back(shape[1] - j - 1);
 274         }
 275     }
 276     EXPECT_TRUE(test::all_close_f(expected_value, actual_value));
 277     EXPECT_EQ(expected_index, actual_index);
 278 }
 279
 280 NGRAPH_TEST(${BACKEND_NAME}, topk_min_sort_value)
 281 {
 282     Shape shape{128, 1000};
 283     Shape rshape{128, 5};
 284     auto A = make_shared<op::Parameter>(element::f32, shape);
 285     auto B = make_shared<op::TopK>(A, 1, element::i32, 5, false, op::TopK::SortType::SORT_VALUES);
 286     auto out_value = B->output(1);
 287     auto out_index = B->output(0);
 288     auto f = make_shared<Function>(OutputVector{out_value, out_index}, ParameterVector{A});
 289
 290     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 291
 292     // Create some tensors for input/output
 293     auto a = backend->create_tensor(element::f32, shape);
 294     vector<float> data;
 295     for (size_t i = 0; i < shape[0]; i++)
 296     {
 297         for (size_t j = 0; j < shape[1]; j++)
 298         {
 299             data.push_back(j);
 300         }
 301     }
 302     copy_data(a, data);
 303
 304     auto result_value = backend->create_tensor(element::f32, rshape);
 305     auto result_index = backend->create_tensor(element::i32, rshape);
 306
 307     auto exec = backend->compile(f);
 308     exec->call({result_value, result_index}, {a});
 309
 310     auto actual_value = read_vector<float>(result_value);
 311     auto actual_index = read_vector<int32_t>(result_index);
 312
 313     for (size_t i = 0; i < rshape[0]; i++)
 314     {
 315         vector<float> expected_value;
 316         vector<int32_t> expected_index;
 317         vector<float> act_value;
 318         vector<int32_t> act_index;
 319         for (size_t j = 0; j < rshape[1]; j++)
 320         {
 321             expected_value.push_back(j);
 322             expected_index.push_back(j);
 323             act_value.push_back(actual_value[rshape[1] * i + j]);
 324             act_index.push_back(actual_index[rshape[1] * i + j]);
 325         }
 326         EXPECT_TRUE(compare_set<float>(expected_value, act_value));
 327         EXPECT_TRUE(compare_set<int32_t>(expected_index, act_index));
 328     }
 329 }
 330
 331 NGRAPH_TEST(${BACKEND_NAME}, topk_max_sort_index)
 332 {
 333     Shape shape{128, 1000};
 334     Shape rshape{128, 5};
 335     auto A = make_shared<op::Parameter>(element::f32, shape);
 336     auto B = make_shared<op::TopK>(A, 1, element::i32, 5, true, op::TopK::SortType::SORT_INDICES);
 337     auto out_value = B->output(1);
 338     auto out_index = B->output(0);
 339     auto f = make_shared<Function>(OutputVector{out_value, out_index}, ParameterVector{A});
 340
 341     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 342
 343     // Create some tensors for input/output
 344     auto a = backend->create_tensor(element::f32, shape);
 345     vector<float> data;
 346     for (size_t i = 0; i < shape[0]; i++)
 347     {
 348         for (size_t j = 0; j < shape[1]; j++)
 349         {
 350             data.push_back(j);
 351         }
 352     }
 353     copy_data(a, data);
 354
 355     auto result_value = backend->create_tensor(element::f32, rshape);
 356     auto result_index = backend->create_tensor(element::i32, rshape);
 357
 358     auto exec = backend->compile(f);
 359     exec->call({result_value, result_index}, {a});
 360
 361     auto actual_value = read_vector<float>(result_value);
 362     auto actual_index = read_vector<int32_t>(result_index);
 363
 364     for (size_t i = 0; i < rshape[0]; i++)
 365     {
 366         vector<float> expected_value;
 367         vector<int32_t> expected_index;
 368         vector<float> act_value;
 369         vector<int32_t> act_index;
 370         for (size_t j = 0; j < rshape[1]; j++)
 371         {
 372             expected_value.push_back(shape[1] - j - 1);
 373             expected_index.push_back(shape[1] - j - 1);
 374             act_value.push_back(actual_value[rshape[1] * i + j]);
 375             act_index.push_back(actual_index[rshape[1] * i + j]);
 376         }
 377         EXPECT_TRUE(compare_set<float>(expected_value, act_value));
 378         EXPECT_TRUE(compare_set<int32_t>(expected_index, act_index));
 379     }
 380 }
 381
 382 NGRAPH_TEST(${BACKEND_NAME}, topk_min_sort_index)
 383 {
 384     Shape shape{128, 1000};
 385     Shape rshape{128, 5};
 386     auto A = make_shared<op::Parameter>(element::f32, shape);
 387     auto B = make_shared<op::TopK>(A, 1, element::i32, 5, false, op::TopK::SortType::SORT_INDICES);
 388     auto out_value = B->output(1);
 389     auto out_index = B->output(0);
 390     auto f = make_shared<Function>(OutputVector{out_value, out_index}, ParameterVector{A});
 391
 392     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 393
 394     // Create some tensors for input/output
 395     auto a = backend->create_tensor(element::f32, shape);
 396     vector<float> data;
 397     for (size_t i = 0; i < shape[0]; i++)
 398     {
 399         for (size_t j = 0; j < shape[1]; j++)
 400         {
 401             data.push_back(j);
 402         }
 403     }
 404     copy_data(a, data);
 405
 406     auto result_value = backend->create_tensor(element::f32, rshape);
 407     auto result_index = backend->create_tensor(element::i32, rshape);
 408
 409     auto exec = backend->compile(f);
 410     exec->call({result_value, result_index}, {a});
 411
 412     auto actual_value = read_vector<float>(result_value);
 413     auto actual_index = read_vector<int32_t>(result_index);
 414
 415     for (size_t i = 0; i < rshape[0]; i++)
 416     {
 417         vector<float> expected_value;
 418         vector<int32_t> expected_index;
 419         vector<float> act_value;
 420         vector<int32_t> act_index;
 421         for (size_t j = 0; j < rshape[1]; j++)
 422         {
 423             expected_value.push_back(j);
 424             expected_index.push_back(j);
 425             act_value.push_back(actual_value[rshape[1] * i + j]);
 426             act_index.push_back(actual_index[rshape[1] * i + j]);
 427         }
 428         EXPECT_TRUE(compare_set<float>(expected_value, act_value));
 429         EXPECT_TRUE(compare_set<int32_t>(expected_index, act_index));
 430     }
 431 }
 432
 433 NGRAPH_TEST(${BACKEND_NAME}, topk_1d_max_all)
 434 {
 435     Shape shape{6};
 436     Shape rshape{6};
 437     auto A = make_shared<op::Parameter>(element::f32, shape);
 438     auto B = make_shared<op::TopK>(A, 0, element::i32, 0, true);
 439     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 440     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 441
 442     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 443
 444     // Create some tensors for input/output
 445     auto a = backend->create_tensor(element::f32, shape);
 446     copy_data(a, vector<float>{1, 2, 3, 4, 5, 6});
 447     auto result0 = backend->create_tensor(element::i32, rshape);
 448     auto result1 = backend->create_tensor(element::f32, rshape);
 449
 450     auto h0 = backend->compile(f0);
 451     h0->call_with_validate({result0}, {a});
 452     EXPECT_EQ((vector<int32_t>{5, 4, 3, 2, 1, 0}), read_vector<int32_t>(result0));
 453     auto h1 = backend->compile(f1);
 454     h1->call_with_validate({result1}, {a});
 455     EXPECT_TRUE(test::all_close_f(
 456         (vector<float>{6, 5, 4, 3, 2, 1}), read_vector<float>(result1), MIN_FLOAT_TOLERANCE_BITS));
 457 }
 458
 459 NGRAPH_TEST(${BACKEND_NAME}, topk_1d_i32_max_all)
 460 {
 461     Shape shape{6};
 462     Shape rshape{6};
 463     auto A = make_shared<op::Parameter>(element::i32, shape);
 464     auto B = make_shared<op::TopK>(A, 0, element::i32, 0, true);
 465     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 466     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 467
 468     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 469
 470     // Create some tensors for input/output
 471     auto a = backend->create_tensor(element::i32, shape);
 472     copy_data(a, vector<int32_t>{1, 2, 3, 4, 5, 6});
 473     auto result0 = backend->create_tensor(element::i32, rshape);
 474     auto result1 = backend->create_tensor(element::i32, rshape);
 475
 476     auto h0 = backend->compile(f0);
 477     h0->call_with_validate({result0}, {a});
 478     EXPECT_EQ((vector<int32_t>{5, 4, 3, 2, 1, 0}), read_vector<int32_t>(result0));
 479     auto h1 = backend->compile(f1);
 480     h1->call_with_validate({result1}, {a});
 481     EXPECT_EQ((vector<int32_t>{6, 5, 4, 3, 2, 1}), read_vector<int32_t>(result1));
 482 }
 483
 484 NGRAPH_TEST(${BACKEND_NAME}, topk_1d_max_partial)
 485 {
 486     Shape shape{6};
 487     Shape rshape{3};
 488     auto A = make_shared<op::Parameter>(element::f32, shape);
 489     auto B = make_shared<op::TopK>(A, 0, element::i32, 3, true);
 490     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 491     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 492
 493     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 494
 495     // Create some tensors for input/output
 496     auto a = backend->create_tensor(element::f32, shape);
 497     copy_data(a, vector<float>{1, 2, 3, 4, 5, 6});
 498     auto result0 = backend->create_tensor(element::i32, rshape);
 499     auto result1 = backend->create_tensor(element::f32, rshape);
 500
 501     auto h0 = backend->compile(f0);
 502     h0->call_with_validate({result0}, {a});
 503     EXPECT_EQ((vector<int32_t>{5, 4, 3}), read_vector<int32_t>(result0));
 504     auto h1 = backend->compile(f1);
 505     h1->call_with_validate({result1}, {a});
 506     EXPECT_TRUE(test::all_close_f(
 507         (vector<float>{6, 5, 4}), read_vector<float>(result1), MIN_FLOAT_TOLERANCE_BITS));
 508 }
 509
 510 NGRAPH_TEST(${BACKEND_NAME}, topk_1d_max_one)
 511 {
 512     Shape shape{6};
 513     Shape rshape{1};
 514     auto A = make_shared<op::Parameter>(element::f32, shape);
 515     auto B = make_shared<op::TopK>(A, 0, element::i32, 1, true);
 516     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 517     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 518
 519     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 520
 521     // Create some tensors for input/output
 522     auto a = backend->create_tensor(element::f32, shape);
 523     copy_data(a, vector<float>{1, 2, 3, 4, 5, 6});
 524     auto result0 = backend->create_tensor(element::i32, rshape);
 525     auto result1 = backend->create_tensor(element::f32, rshape);
 526
 527     auto h0 = backend->compile(f0);
 528     h0->call_with_validate({result0}, {a});
 529     EXPECT_EQ((vector<int32_t>{5}), read_vector<int32_t>(result0));
 530     auto h1 = backend->compile(f1);
 531     h1->call_with_validate({result1}, {a});
 532     EXPECT_TRUE(test::all_close_f(
 533         (vector<float>{6}), read_vector<float>(result1), MIN_FLOAT_TOLERANCE_BITS));
 534 }
 535
 536 NGRAPH_TEST(${BACKEND_NAME}, topk_1d_min_all)
 537 {
 538     Shape shape{6};
 539     Shape rshape{6};
 540     auto A = make_shared<op::Parameter>(element::f32, shape);
 541     auto B = make_shared<op::TopK>(A, 0, element::i32, 0, false);
 542     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 543     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 544
 545     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 546
 547     // Create some tensors for input/output
 548     auto a = backend->create_tensor(element::f32, shape);
 549     copy_data(a, vector<float>{6, 5, 4, 3, 2, 1});
 550     auto result0 = backend->create_tensor(element::i32, rshape);
 551     auto result1 = backend->create_tensor(element::f32, rshape);
 552
 553     auto h0 = backend->compile(f0);
 554     h0->call_with_validate({result0}, {a});
 555     EXPECT_EQ((vector<int32_t>{5, 4, 3, 2, 1, 0}), read_vector<int32_t>(result0));
 556     auto h1 = backend->compile(f1);
 557     h1->call_with_validate({result1}, {a});
 558     EXPECT_TRUE(test::all_close_f(
 559         (vector<float>{1, 2, 3, 4, 5, 6}), read_vector<float>(result1), MIN_FLOAT_TOLERANCE_BITS));
 560 }
 561
 562 NGRAPH_TEST(${BACKEND_NAME}, topk_1d_min_partial)
 563 {
 564     Shape shape{6};
 565     Shape rshape{3};
 566     auto A = make_shared<op::Parameter>(element::f32, shape);
 567     auto B = make_shared<op::TopK>(A, 0, element::i32, 3, false);
 568     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 569     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 570
 571     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 572
 573     // Create some tensors for input/output
 574     auto a = backend->create_tensor(element::f32, shape);
 575     copy_data(a, vector<float>{6, 5, 4, 3, 2, 1});
 576     auto result0 = backend->create_tensor(element::i32, rshape);
 577     auto result1 = backend->create_tensor(element::f32, rshape);
 578
 579     auto h0 = backend->compile(f0);
 580     h0->call_with_validate({result0}, {a});
 581     EXPECT_EQ((vector<int32_t>{5, 4, 3}), read_vector<int32_t>(result0));
 582     auto h1 = backend->compile(f1);
 583     h1->call_with_validate({result1}, {a});
 584     EXPECT_TRUE(test::all_close_f(
 585         (vector<float>{1, 2, 3}), read_vector<float>(result1), MIN_FLOAT_TOLERANCE_BITS));
 586 }
 587
 588 NGRAPH_TEST(${BACKEND_NAME}, topk_1d_min_one)
 589 {
 590     Shape shape{6};
 591     Shape rshape{1};
 592     auto A = make_shared<op::Parameter>(element::f32, shape);
 593     auto B = make_shared<op::TopK>(A, 0, element::i32, 1, false);
 594     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 595     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 596
 597     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 598
 599     // Create some tensors for input/output
 600     auto a = backend->create_tensor(element::f32, shape);
 601     copy_data(a, vector<float>{6, 5, 4, 3, 2, 1});
 602     auto result0 = backend->create_tensor(element::i32, rshape);
 603     auto result1 = backend->create_tensor(element::f32, rshape);
 604
 605     auto h0 = backend->compile(f0);
 606     h0->call_with_validate({result0}, {a});
 607     EXPECT_EQ((vector<int32_t>{5}), read_vector<int32_t>(result0));
 608     auto h1 = backend->compile(f1);
 609     h1->call_with_validate({result1}, {a});
 610     EXPECT_TRUE(test::all_close_f(
 611         (vector<float>{1}), read_vector<float>(result1), MIN_FLOAT_TOLERANCE_BITS));
 612 }
 613
 614 NGRAPH_TEST(${BACKEND_NAME}, topk_3d_max_all)
 615 {
 616     Shape shape{2, 3, 2};
 617     Shape rshape{2, 3, 2};
 618     auto A = make_shared<op::Parameter>(element::f32, shape);
 619     auto B = make_shared<op::TopK>(A, 1, element::i32, 0, true);
 620     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 621     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 622
 623     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 624
 625     // Create some tensors for input/output
 626     auto a = backend->create_tensor(element::f32, shape);
 627     copy_data(a, vector<float>{9, 2, 10, 12, 8, 4, 6, 1, 5, 3, 11, 7});
 628     auto result0 = backend->create_tensor(element::i32, rshape);
 629     auto result1 = backend->create_tensor(element::f32, rshape);
 630
 631     auto h0 = backend->compile(f0);
 632     h0->call_with_validate({result0}, {a});
 633     EXPECT_EQ((vector<int32_t>{1, 1, 0, 2, 2, 0, 2, 2, 0, 1, 1, 0}), read_vector<int32_t>(result0));
 634     auto h1 = backend->compile(f1);
 635     h1->call_with_validate({result1}, {a});
 636     EXPECT_TRUE(test::all_close_f((vector<float>{10, 12, 9, 4, 8, 2, 11, 7, 6, 3, 5, 1}),
 637                                   read_vector<float>(result1),
 638                                   MIN_FLOAT_TOLERANCE_BITS));
 639 }
 640
 641 NGRAPH_TEST(${BACKEND_NAME}, topk_int64)
 642 {
 643     Shape shape{2, 3, 2};
 644     Shape rshape{2, 3, 2};
 645     auto A = make_shared<op::Parameter>(element::f32, shape);
 646     auto B = make_shared<op::TopK>(A, 1, element::i64, 0, true);
 647     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 648     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 649
 650     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 651
 652     // Create some tensors for input/output
 653     auto a = backend->create_tensor(element::f32, shape);
 654     copy_data(a, vector<float>{9, 2, 10, 12, 8, 4, 6, 1, 5, 3, 11, 7});
 655     auto result0 = backend->create_tensor(element::i64, rshape);
 656     auto result1 = backend->create_tensor(element::f32, rshape);
 657
 658     auto h0 = backend->compile(f0);
 659     h0->call_with_validate({result0}, {a});
 660     EXPECT_EQ((vector<int64_t>{1, 1, 0, 2, 2, 0, 2, 2, 0, 1, 1, 0}), read_vector<int64_t>(result0));
 661     auto h1 = backend->compile(f1);
 662     h1->call_with_validate({result1}, {a});
 663     EXPECT_TRUE(test::all_close_f((vector<float>{10, 12, 9, 4, 8, 2, 11, 7, 6, 3, 5, 1}),
 664                                   read_vector<float>(result1),
 665                                   MIN_FLOAT_TOLERANCE_BITS));
 666 }
 667
 668 NGRAPH_TEST(${BACKEND_NAME}, topk_5d_max_partial)
 669 {
 670     Shape shape{2, 6, 3, 2, 4};
 671     Shape rshape{2, 2, 3, 2, 4};
 672     auto A = make_shared<op::Parameter>(element::f32, shape);
 673     auto B = make_shared<op::TopK>(A, 1, element::i32, 2, true);
 674     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 675     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 676
 677     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 678
 679     // Create some tensors for input/output
 680     auto a = backend->create_tensor(element::f32, shape);
 681     copy_data(
 682         a,
 683         vector<float>{
 684             1.,   73.,  9.,   81.,  17.,  89.,  2.,   74.,  10.,  82.,  18.,  90.,  3.,   75.,
 685             11.,  83.,  19.,  91.,  4.,   76.,  12.,  84.,  20.,  92.,  145., 217., 153., 225.,
 686             161., 233., 146., 218., 154., 226., 162., 234., 147., 219., 155., 227., 163., 235.,
 687             148., 220., 156., 228., 164., 236., 5.,   77.,  13.,  85.,  21.,  93.,  6.,   78.,
 688             14.,  86.,  22.,  94.,  7.,   79.,  15.,  87.,  23.,  95.,  8.,   80.,  16.,  88.,
 689             24.,  96.,  149., 221., 157., 229., 165., 27.,  150., 222., 158., 230., 166., 23.,
 690             151., 223., 159., 231., 17.,  39.,  2.,   224., 160., 232., 168., 240., 25.,  97.,
 691             33.,  105., 41.,  113., 26.,  98.,  34.,  106., 42.,  114., 27.,  99.,  35.,  107.,
 692             43.,  115., 28.,  100., 36.,  108., 44.,  116., 169., 241., 177., 249., 185., 25.,
 693             170., 242., 178., 250., 186., 258., 171., 243., 179., 251., 187., 259., 172., 24.,
 694             180., 252., 188., 260., 29.,  101., 37.,  109., 45.,  117., 30.,  102., 38.,  10.,
 695             46.,  118., 31.,  103., 39.,  111., 47.,  119., 32.,  104., 40.,  112., 48.,  20.,
 696             173., 245., 181., 253., 189., 261., 174., 246., 182., 254., 190., 262., 175., 27.,
 697             183., 255., 191., 263., 176., 248., 184., 256., 192., 264., 49.,  121., 57.,  129.,
 698             65.,  137., 50.,  122., 58.,  130., 66.,  138., 51.,  123., 59.,  131., 67.,  139.,
 699             52.,  124., 60.,  132., 68.,  140., 193., 265., 201., 273., 209., 281., 194., 266.,
 700             202., 274., 210., 43.,  115., 28.,  100., 36.,  108., 44.,  116., 169., 241., 177.,
 701             212., 284., 53.,  125., 61.,  133., 69.,  141., 54.,  126., 62.,  134., 70.,  142.,
 702             55.,  127., 63.,  135., 71.,  143., 56.,  128., 64.,  136., 72.,  144., 197., 269.,
 703             205., 277., 213., 285., 198., 270., 206., 278., 214., 286., 199., 271., 207., 279.,
 704             215., 287., 200., 272., 208., 280., 216., 288.});
 705
 706     auto result0 = backend->create_tensor(element::i32, rshape);
 707     auto result1 = backend->create_tensor(element::f32, rshape);
 708
 709     auto h0 = backend->compile(f0);
 710     h0->call_with_validate({result0}, {a});
 711     EXPECT_EQ(
 712         (vector<int32_t>{5, 5, 5, 5, 5, 1, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 3, 5, 5, 5, 5,
 713                          3, 3, 3, 3, 3, 4, 3, 3, 3, 3, 3, 1, 3, 3, 3, 3, 1, 1, 1, 1, 3, 3, 3, 3,
 714                          5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 3, 5, 5, 5,
 715                          3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 1, 1, 4, 1, 1, 1, 1, 1, 1, 5, 1, 3, 3}),
 716         read_vector<int32_t>(result0));
 717
 718     auto h1 = backend->compile(f1);
 719     h1->call_with_validate({result1}, {a});
 720     EXPECT_TRUE(test::all_close_f(
 721         (vector<float>{169, 241, 177, 249, 185, 233, 170, 242, 178, 250, 186, 258, 171, 243,
 722                        179, 251, 187, 259, 172, 224, 180, 252, 188, 260, 149, 221, 157, 229,
 723                        165, 113, 150, 222, 158, 230, 166, 234, 151, 223, 159, 231, 163, 235,
 724                        148, 220, 160, 232, 168, 240, 197, 269, 205, 277, 213, 285, 198, 270,
 725                        206, 278, 214, 286, 199, 271, 207, 279, 215, 287, 200, 272, 241, 280,
 726                        216, 288, 193, 265, 201, 273, 209, 281, 194, 266, 202, 274, 210, 262,
 727                        175, 127, 183, 255, 191, 263, 176, 248, 208, 256, 212, 284}),
 728         read_vector<float>(result1),
 729         MIN_FLOAT_TOLERANCE_BITS));
 730 }
 731
 732 NGRAPH_TEST(${BACKEND_NAME}, topk_3d_max_partial)
 733 {
 734     Shape shape{2, 3, 2};
 735     Shape rshape{2, 2, 2};
 736     auto A = make_shared<op::Parameter>(element::f32, shape);
 737     auto B = make_shared<op::TopK>(A, 1, element::i32, 2, true);
 738     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 739     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 740
 741     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 742
 743     // Create some tensors for input/output
 744     auto a = backend->create_tensor(element::f32, shape);
 745     copy_data(a, vector<float>{9, 2, 10, 12, 8, 4, 6, 1, 5, 3, 11, 7});
 746     auto result0 = backend->create_tensor(element::i32, rshape);
 747     auto result1 = backend->create_tensor(element::f32, rshape);
 748
 749     auto h0 = backend->compile(f0);
 750     h0->call_with_validate({result0}, {a});
 751     EXPECT_EQ((vector<int32_t>{1, 1, 0, 2, 2, 2, 0, 1}), read_vector<int32_t>(result0));
 752     auto h1 = backend->compile(f1);
 753     h1->call_with_validate({result1}, {a});
 754     EXPECT_TRUE(test::all_close_f((vector<float>{10, 12, 9, 4, 11, 7, 6, 3}),
 755                                   read_vector<float>(result1),
 756                                   MIN_FLOAT_TOLERANCE_BITS));
 757 }
 758
 759 NGRAPH_TEST(${BACKEND_NAME}, topk_3d_max_one)
 760 {
 761     Shape shape{2, 3, 2};
 762     Shape rshape{2, 1, 2};
 763     auto A = make_shared<op::Parameter>(element::f32, shape);
 764     auto B = make_shared<op::TopK>(A, 1, element::i32, 1, true);
 765     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 766     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 767
 768     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 769
 770     // Create some tensors for input/output
 771     auto a = backend->create_tensor(element::f32, shape);
 772     copy_data(a, vector<float>{9, 2, 10, 12, 8, 4, 6, 1, 5, 3, 11, 7});
 773     auto result0 = backend->create_tensor(element::i32, rshape);
 774     auto result1 = backend->create_tensor(element::f32, rshape);
 775
 776     auto h0 = backend->compile(f0);
 777     h0->call_with_validate({result0}, {a});
 778     EXPECT_EQ((vector<int32_t>{1, 1, 2, 2}), read_vector<int32_t>(result0));
 779     auto h1 = backend->compile(f1);
 780     h1->call_with_validate({result1}, {a});
 781     EXPECT_TRUE(test::all_close_f(
 782         (vector<float>{10, 12, 11, 7}), read_vector<float>(result1), MIN_FLOAT_TOLERANCE_BITS));
 783 }
 784
 785 NGRAPH_TEST(${BACKEND_NAME}, topk_3d_min_all)
 786 {
 787     Shape shape{2, 3, 2};
 788     Shape rshape{2, 3, 2};
 789     auto A = make_shared<op::Parameter>(element::f32, shape);
 790     auto B = make_shared<op::TopK>(A, 1, element::i32, 0, false);
 791     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 792     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 793
 794     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 795
 796     // Create some tensors for input/output
 797     auto a = backend->create_tensor(element::f32, shape);
 798     copy_data(a, vector<float>{12, 2, 10, 9, 8, 4, 6, 1, 5, 3, 11, 7});
 799     auto result0 = backend->create_tensor(element::i32, rshape);
 800     auto result1 = backend->create_tensor(element::f32, rshape);
 801
 802     auto h0 = backend->compile(f0);
 803     h0->call_with_validate({result0}, {a});
 804     EXPECT_EQ((vector<int32_t>{2, 0, 1, 2, 0, 1, 1, 0, 0, 1, 2, 2}), read_vector<int32_t>(result0));
 805     auto h1 = backend->compile(f1);
 806     h1->call_with_validate({result1}, {a});
 807     EXPECT_TRUE(test::all_close_f((vector<float>{8, 2, 10, 4, 12, 9, 5, 1, 6, 3, 11, 7}),
 808                                   read_vector<float>(result1),
 809                                   MIN_FLOAT_TOLERANCE_BITS));
 810 }
 811
 812 NGRAPH_TEST(${BACKEND_NAME}, topk_3d_min_partial)
 813 {
 814     Shape shape{2, 3, 2};
 815     Shape rshape{2, 2, 2};
 816     auto A = make_shared<op::Parameter>(element::f32, shape);
 817     auto B = make_shared<op::TopK>(A, 1, element::i32, 2, false);
 818     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 819     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 820
 821     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 822
 823     // Create some tensors for input/output
 824     auto a = backend->create_tensor(element::f32, shape);
 825     copy_data(a, vector<float>{12, 2, 10, 9, 8, 4, 6, 1, 5, 3, 11, 7});
 826     auto result0 = backend->create_tensor(element::i32, rshape);
 827     auto result1 = backend->create_tensor(element::f32, rshape);
 828
 829     auto h0 = backend->compile(f0);
 830     h0->call_with_validate({result0}, {a});
 831     EXPECT_EQ((vector<int32_t>{2, 0, 1, 2, 1, 0, 0, 1}), read_vector<int32_t>(result0));
 832     auto h1 = backend->compile(f1);
 833     h1->call_with_validate({result1}, {a});
 834     EXPECT_TRUE(test::all_close_f((vector<float>{8, 2, 10, 4, 5, 1, 6, 3}),
 835                                   read_vector<float>(result1),
 836                                   MIN_FLOAT_TOLERANCE_BITS));
 837 }
 838
 839 NGRAPH_TEST(${BACKEND_NAME}, topk_3d_min_one)
 840 {
 841     Shape shape{2, 3, 2};
 842     Shape rshape{2, 1, 2};
 843     auto A = make_shared<op::Parameter>(element::f32, shape);
 844     auto B = make_shared<op::TopK>(A, 1, element::i32, 1, false);
 845     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 846     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 847
 848     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 849
 850     // Create some tensors for input/output
 851     auto a = backend->create_tensor(element::f32, shape);
 852     copy_data(a, vector<float>{12, 2, 10, 9, 8, 4, 6, 1, 5, 3, 11, 7});
 853     auto result0 = backend->create_tensor(element::i32, rshape);
 854     auto result1 = backend->create_tensor(element::f32, rshape);
 855
 856     auto h0 = backend->compile(f0);
 857     h0->call_with_validate({result0}, {a});
 858     EXPECT_EQ((vector<int32_t>{2, 0, 1, 0}), read_vector<int32_t>(result0));
 859     auto h1 = backend->compile(f1);
 860     h1->call_with_validate({result1}, {a});
 861     EXPECT_TRUE(test::all_close_f(
 862         (vector<float>{8, 2, 5, 1}), read_vector<float>(result1), MIN_FLOAT_TOLERANCE_BITS));
 863 }
 864
 865 NGRAPH_TEST(${BACKEND_NAME}, topk_2d_max_all)
 866 {
 867     Shape shape{4, 3};
 868     Shape rshape{4, 3};
 869     auto A = make_shared<op::Parameter>(element::f32, shape);
 870     auto B = make_shared<op::TopK>(A, 0, element::i32, 4, true);
 871     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 872     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 873
 874     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 875
 876     // Create some tensors for input/output
 877     auto a = backend->create_tensor(element::f32, shape);
 878     copy_data(a, vector<float>{9, 2, 10, 12, 8, 4, 6, 1, 5, 3, 11, 7});
 879     auto result0 = backend->create_tensor(element::i32, rshape);
 880     auto result1 = backend->create_tensor(element::f32, rshape);
 881
 882     auto h0 = backend->compile(f0);
 883     h0->call_with_validate({result0}, {a});
 884     EXPECT_EQ((vector<int32_t>{1, 3, 0, 0, 1, 3, 2, 0, 2, 3, 2, 1}), read_vector<int32_t>(result0));
 885     auto h1 = backend->compile(f1);
 886     h1->call_with_validate({result1}, {a});
 887     EXPECT_TRUE(test::all_close_f((vector<float>{12, 11, 10, 9, 8, 7, 6, 2, 5, 3, 1, 4}),
 888                                   read_vector<float>(result1),
 889                                   MIN_FLOAT_TOLERANCE_BITS));
 890 }
 891
 892 NGRAPH_TEST(${BACKEND_NAME}, topk_2d_max_partial)
 893 {
 894     Shape shape{4, 3};
 895     Shape rshape{2, 3};
 896     auto A = make_shared<op::Parameter>(element::f32, shape);
 897     auto B = make_shared<op::TopK>(A, 0, element::i32, 2, true);
 898     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 899     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 900
 901     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 902
 903     // Create some tensors for input/output
 904     auto a = backend->create_tensor(element::f32, shape);
 905     copy_data(a, vector<float>{9, 2, 10, 12, 8, 4, 6, 1, 5, 3, 11, 7});
 906     auto result0 = backend->create_tensor(element::i32, rshape);
 907     auto result1 = backend->create_tensor(element::f32, rshape);
 908
 909     auto h0 = backend->compile(f0);
 910     h0->call_with_validate({result0}, {a});
 911     EXPECT_EQ((vector<int32_t>{1, 3, 0, 0, 1, 3}), read_vector<int32_t>(result0));
 912     auto h1 = backend->compile(f1);
 913     h1->call_with_validate({result1}, {a});
 914     EXPECT_TRUE(test::all_close_f((vector<float>{12, 11, 10, 9, 8, 7}),
 915                                   read_vector<float>(result1),
 916                                   MIN_FLOAT_TOLERANCE_BITS));
 917 }
 918
 919 NGRAPH_TEST(${BACKEND_NAME}, topk_2d_max_one)
 920 {
 921     Shape shape{4, 3};
 922     Shape rshape{1, 3};
 923     auto A = make_shared<op::Parameter>(element::f32, shape);
 924     auto B = make_shared<op::TopK>(A, 0, element::i32, 1, true);
 925     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 926     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 927
 928     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 929
 930     // Create some tensors for input/output
 931     auto a = backend->create_tensor(element::f32, shape);
 932     copy_data(a, vector<float>{9, 2, 10, 12, 8, 4, 6, 1, 5, 3, 11, 7});
 933     auto result0 = backend->create_tensor(element::i32, rshape);
 934     auto result1 = backend->create_tensor(element::f32, rshape);
 935
 936     auto h0 = backend->compile(f0);
 937     h0->call_with_validate({result0}, {a});
 938     EXPECT_EQ((vector<int32_t>{1, 3, 0}), read_vector<int32_t>(result0));
 939     auto h1 = backend->compile(f1);
 940     h1->call_with_validate({result1}, {a});
 941     EXPECT_TRUE(test::all_close_f(
 942         (vector<float>{12, 11, 10}), read_vector<float>(result1), MIN_FLOAT_TOLERANCE_BITS));
 943 }
 944
 945 NGRAPH_TEST(${BACKEND_NAME}, topk_2d_max_one_with_equal_values)
 946 {
 947     Shape shape{2, 4};
 948     Shape rshape{2, 1};
 949     auto A = make_shared<op::Parameter>(element::f32, shape);
 950     auto B = make_shared<op::TopK>(A, 1, element::i32, 1, true);
 951     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 952     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 953
 954     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 955
 956     // Create some tensors for input/output
 957     auto a = backend->create_tensor(element::f32, shape);
 958     copy_data(a, vector<float>{1, 3, 2, 4, 1, 3, 3, 2});
 959     auto result0 = backend->create_tensor(element::i32, rshape);
 960     auto result1 = backend->create_tensor(element::f32, rshape);
 961
 962     auto h0 = backend->compile(f0);
 963     h0->call_with_validate({result0}, {a});
 964     EXPECT_EQ((vector<int32_t>{3, 1}), read_vector<int32_t>(result0));
 965     auto h1 = backend->compile(f1);
 966     h1->call_with_validate({result1}, {a});
 967     EXPECT_TRUE(test::all_close_f(
 968         (vector<float>{4, 3}), read_vector<float>(result1), MIN_FLOAT_TOLERANCE_BITS));
 969 }
 970
 971 NGRAPH_TEST(${BACKEND_NAME}, topk_2d_min_all)
 972 {
 973     Shape shape{4, 3};
 974     Shape rshape{4, 3};
 975     auto A = make_shared<op::Parameter>(element::f32, shape);
 976     auto B = make_shared<op::TopK>(A, 0, element::i32, 4, false);
 977     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
 978     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
 979
 980     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 981
 982     // Create some tensors for input/output
 983     auto a = backend->create_tensor(element::f32, shape);
 984     copy_data(a, vector<float>{12, 2, 10, 9, 8, 4, 6, 1, 5, 3, 11, 7});
 985     auto result0 = backend->create_tensor(element::i32, rshape);
 986     auto result1 = backend->create_tensor(element::f32, rshape);
 987
 988     auto h0 = backend->compile(f0);
 989     h0->call_with_validate({result0}, {a});
 990     EXPECT_EQ((vector<int32_t>{3, 2, 1, 2, 0, 2, 1, 1, 3, 0, 3, 0}), read_vector<int32_t>(result0));
 991     auto h1 = backend->compile(f1);
 992     h1->call_with_validate({result1}, {a});
 993     EXPECT_TRUE(test::all_close_f((vector<float>{3, 1, 4, 6, 2, 5, 9, 8, 7, 12, 11, 10}),
 994                                   read_vector<float>(result1),
 995                                   MIN_FLOAT_TOLERANCE_BITS));
 996 }
 997
 998 NGRAPH_TEST(${BACKEND_NAME}, topk_2d_min_partial)
 999 {
1000     Shape shape{4, 3};
1001     Shape rshape{2, 3};
1002     auto A = make_shared<op::Parameter>(element::f32, shape);
1003     auto B = make_shared<op::TopK>(A, 0, element::i32, 2, false);
1004     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
1005     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
1006
1007     auto backend = runtime::Backend::create("${BACKEND_NAME}");
1008
1009     // Create some tensors for input/output
1010     auto a = backend->create_tensor(element::f32, shape);
1011     copy_data(a, vector<float>{12, 2, 10, 9, 8, 4, 6, 1, 5, 3, 11, 7});
1012     auto result0 = backend->create_tensor(element::i32, rshape);
1013     auto result1 = backend->create_tensor(element::f32, rshape);
1014
1015     auto h0 = backend->compile(f0);
1016     h0->call_with_validate({result0}, {a});
1017     EXPECT_EQ((vector<int32_t>{3, 2, 1, 2, 0, 2}), read_vector<int32_t>(result0));
1018     auto h1 = backend->compile(f1);
1019     h1->call_with_validate({result1}, {a});
1020     EXPECT_TRUE(test::all_close_f(
1021         (vector<float>{3, 1, 4, 6, 2, 5}), read_vector<float>(result1), MIN_FLOAT_TOLERANCE_BITS));
1022 }
1023
1024 NGRAPH_TEST(${BACKEND_NAME}, topk_2d_min_one)
1025 {
1026     Shape shape{4, 3};
1027     Shape rshape{1, 3};
1028     auto A = make_shared<op::Parameter>(element::f32, shape);
1029     auto B = make_shared<op::TopK>(A, 0, element::i32, 1, false);
1030     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
1031     auto f1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
1032
1033     auto backend = runtime::Backend::create("${BACKEND_NAME}");
1034
1035     // Create some tensors for input/output
1036     auto a = backend->create_tensor(element::f32, shape);
1037     copy_data(a, vector<float>{12, 2, 10, 9, 8, 4, 6, 1, 5, 3, 11, 7});
1038     auto result0 = backend->create_tensor(element::i32, rshape);
1039     auto result1 = backend->create_tensor(element::f32, rshape);
1040
1041     auto h0 = backend->compile(f0);
1042     h0->call_with_validate({result0}, {a});
1043     EXPECT_EQ((vector<int32_t>{3, 2, 1}), read_vector<int32_t>(result0));
1044     auto h1 = backend->compile(f1);
1045     h1->call_with_validate({result1}, {a});
1046     EXPECT_TRUE(test::all_close_f(
1047         (vector<float>{3, 1, 4}), read_vector<float>(result1), MIN_FLOAT_TOLERANCE_BITS));
1048 }
1049
1050 NGRAPH_TEST(${BACKEND_NAME}, topk_3d_large_input_max)
1051 {
1052     Shape shape{4, 8192, 5};
1053     auto A = make_shared<op::Parameter>(element::f32, shape);
1054
1055     auto B = make_shared<op::TopK>(A, 1, element::i32, 10, true);
1056
1057     auto interp_f_0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
1058     auto interp_f_1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
1059     auto gpu_f_0 = ngraph::clone_function(*interp_f_0);
1060     auto gpu_f_1 = ngraph::clone_function(*interp_f_1);
1061
1062     vector<vector<float>> args;
1063     for (shared_ptr<op::Parameter> param : interp_f_0->get_parameters())
1064     {
1065         vector<float> tensor_val(shape_size(param->get_shape()));
1066         iota(tensor_val.begin(), tensor_val.end(), 0.0f);
1067         args.push_back(tensor_val);
1068     }
1069
1070     auto interp_results_0 = execute<float, int32_t>(interp_f_0, args, "INTERPRETER");
1071     auto gpu_results_0 = execute<float, int32_t>(gpu_f_0, args, "${BACKEND_NAME}");
1072     for (size_t i = 0; i < gpu_results_0.size(); i++)
1073     {
1074         EXPECT_EQ(gpu_results_0.at(i), interp_results_0.at(i));
1075     }
1076
1077     auto interp_results_1 = execute(interp_f_1, args, "INTERPRETER");
1078     auto gpu_results_1 = execute(gpu_f_1, args, "${BACKEND_NAME}");
1079
1080     for (size_t i = 0; i < gpu_results_1.size(); i++)
1081     {
1082         EXPECT_TRUE(test::all_close_f(
1083             gpu_results_1.at(i), interp_results_1.at(i), MIN_FLOAT_TOLERANCE_BITS));
1084     }
1085 }
1086
1087 NGRAPH_TEST(${BACKEND_NAME}, topk_3d_large_input_min)
1088 {
1089     Shape shape{4, 8192, 5};
1090     auto A = make_shared<op::Parameter>(element::f32, shape);
1091
1092     auto B = make_shared<op::TopK>(A, 1, element::i32, 10, false);
1093
1094     auto interp_f_0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
1095     auto interp_f_1 = make_shared<Function>(OutputVector{B->output(1)}, ParameterVector{A});
1096     auto gpu_f_0 = ngraph::clone_function(*interp_f_0);
1097     auto gpu_f_1 = ngraph::clone_function(*interp_f_1);
1098
1099     vector<vector<float>> args;
1100     for (shared_ptr<op::Parameter> param : interp_f_0->get_parameters())
1101     {
1102         vector<float> tensor_val(shape_size(param->get_shape()));
1103         iota(tensor_val.begin(), tensor_val.end(), 0.0f);
1104         args.push_back(tensor_val);
1105     }
1106
1107     auto interp_results_0 = execute<float, int32_t>(interp_f_0, args, "INTERPRETER");
1108     auto gpu_results_0 = execute<float, int32_t>(gpu_f_0, args, "${BACKEND_NAME}");
1109     for (size_t i = 0; i < gpu_results_0.size(); i++)
1110     {
1111         EXPECT_EQ(gpu_results_0.at(i), interp_results_0.at(i));
1112     }
1113
1114     auto interp_results_1 = execute(interp_f_1, args, "INTERPRETER");
1115     auto gpu_results_1 = execute(gpu_f_1, args, "${BACKEND_NAME}");
1116
1117     for (size_t i = 0; i < gpu_results_1.size(); i++)
1118     {
1119         EXPECT_TRUE(test::all_close_f(
1120             gpu_results_1.at(i), interp_results_1.at(i), MIN_FLOAT_TOLERANCE_BITS));
1121     }
1122 }
1123
1124 NGRAPH_TEST(${BACKEND_NAME}, topk_3d_single_output)
1125 {
1126     Shape shape{2, 3, 2};
1127     Shape rshape{2, 2, 2};
1128     auto A = make_shared<op::Parameter>(element::f32, shape);
1129     auto B = make_shared<op::TopK>(A, 1, element::i32, 2, false);
1130     auto f0 = make_shared<Function>(OutputVector{B->output(0)}, ParameterVector{A});
1131
1132     auto backend = runtime::Backend::create("${BACKEND_NAME}");
1133
1134     // Create some tensors for input/output
1135     auto a = backend->create_tensor(element::f32, shape);
1136     copy_data(a, vector<float>{12, 2, 10, 9, 8, 4, 6, 1, 5, 3, 11, 7});
1137     auto result0 = backend->create_tensor(element::i32, rshape);
1138
1139     auto h0 = backend->compile(f0);
1140     h0->call_with_validate({result0}, {a});
1141     EXPECT_EQ((vector<int32_t>{2, 0, 1, 2, 1, 0, 0, 1}), read_vector<int32_t>(result0));
1142 }
1143
1144 NGRAPH_TEST(${BACKEND_NAME}, topk_v1_invalid_strings)
1145 {
1146     const auto data = make_shared<op::Parameter>(element::f32, Shape{1, 2, 3});
1147     const auto k = op::Constant::create(element::i64, Shape{}, {1});
1148     EXPECT_THROW(op::v1::TopK(data, k, 0, "max", "invalid_mode"), ngraph::CheckFailure);
1149     EXPECT_THROW(op::v1::TopK(data, k, 0, "invalid_sort", "index"), ngraph::CheckFailure);
1150 }
1151
1152 NGRAPH_TEST(${BACKEND_NAME}, topk_v1_invalid_k)
1153 {
1154     const auto data = make_shared<op::Parameter>(element::f32, Shape{1, 2, 3});
1155
1156     // K must be a scalar
1157     const auto k_non_scalar = op::Constant::create(element::i64, Shape{2}, {1, 2});
1158     EXPECT_THROW(op::v1::TopK(data, k_non_scalar, 0, "max", "index"),
1159                  ngraph::NodeValidationFailure);
1160
1161     // K can only be i8, i32 or i64
1162     const auto k_float = op::Constant::create(element::f32, Shape{}, {1.0f});
1163     EXPECT_THROW(op::v1::TopK(data, k_float, 0, "max", "index"), ngraph::NodeValidationFailure);
1164
1165     // the value of K must be positive
1166     const auto k_negative = op::Constant::create(element::i8, Shape{}, {-1});
1167     EXPECT_THROW(op::v1::TopK(data, k_negative, 0, "max", "index"), ngraph::NodeValidationFailure);
1168 }