From 2ac35247ea13de658af255685e3440f9b033eb96 Mon Sep 17 00:00:00 2001
From: Mateusz Tabaka <mateusz.tabaka@intel.com>
Date: Mon, 27 Jul 2020 12:40:27 +0200
Subject: [PATCH] Add tests for ArgMin/ArgMax with float inputs (#1429)

---
 ngraph/test/models/onnx/argmax_float.prototxt | 55 +++++++++++++++++++
 ngraph/test/models/onnx/argmin_float.prototxt | 55 +++++++++++++++++++
 ngraph/test/onnx/onnx_import.in.cpp           | 22 ++++++++
 ngraph/test/runtime/ie/unit_test.manifest     |  4 ++
 4 files changed, 136 insertions(+)
 create mode 100644 ngraph/test/models/onnx/argmax_float.prototxt
 create mode 100644 ngraph/test/models/onnx/argmin_float.prototxt

diff --git a/ngraph/test/models/onnx/argmax_float.prototxt b/ngraph/test/models/onnx/argmax_float.prototxt
new file mode 100644
index 00000000000..755c50af476
--- /dev/null
+++ b/ngraph/test/models/onnx/argmax_float.prototxt
@@ -0,0 +1,55 @@
+ir_version: 3
+producer_name: "nGraph ONNX Importer"
+graph {
+  node {
+    input: "data"
+    output: "reduced"
+    name: "node1"
+    op_type: "ArgMax"
+    attribute {
+      name: "keepdims"
+      i: 1
+      type: INT
+    }
+    attribute {
+      name: "axis"
+      i: 0
+      type: INT
+    }
+    doc_string: "ArgMax"
+    domain: ""
+  }
+  name: "test"
+  input {
+    name: "data"
+    type {
+      tensor_type {
+        elem_type: 1
+        shape {
+          dim {
+            dim_value: 4
+          }
+          dim {
+            dim_value: 3
+          }
+        }
+      }
+    }
+  }
+  output {
+    name: "reduced"
+    type {
+      tensor_type {
+        elem_type: 7
+        shape {
+          dim {
+            dim_value: 3
+          }
+        }
+      }
+    }
+  }
+}
+opset_import {
+  version: 7
+}
diff --git a/ngraph/test/models/onnx/argmin_float.prototxt b/ngraph/test/models/onnx/argmin_float.prototxt
new file mode 100644
index 00000000000..6e8f9703159
--- /dev/null
+++ b/ngraph/test/models/onnx/argmin_float.prototxt
@@ -0,0 +1,55 @@
+ir_version: 3
+producer_name: "nGraph ONNX Importer"
+graph {
+  node {
+    input: "data"
+    output: "reduced"
+    name: "node1"
+    op_type: "ArgMin"
+    attribute {
+      name: "keepdims"
+      i: 0
+      type: INT
+    }
+    attribute {
+      name: "axis"
+      i: 1
+      type: INT
+    }
+    doc_string: "ArgMin"
+    domain: ""
+  }
+  name: "test"
+  input {
+    name: "data"
+    type {
+      tensor_type {
+        elem_type: 1
+        shape {
+          dim {
+            dim_value: 4
+          }
+          dim {
+            dim_value: 3
+          }
+        }
+      }
+    }
+  }
+  output {
+    name: "reduced"
+    type {
+      tensor_type {
+        elem_type: 7
+        shape {
+          dim {
+            dim_value: 4
+          }
+        }
+      }
+    }
+  }
+}
+opset_import {
+  version: 7
+}
diff --git a/ngraph/test/onnx/onnx_import.in.cpp b/ngraph/test/onnx/onnx_import.in.cpp
index 9f519f70921..6d475608f82 100644
--- a/ngraph/test/onnx/onnx_import.in.cpp
+++ b/ngraph/test/onnx/onnx_import.in.cpp
@@ -1626,6 +1626,28 @@ NGRAPH_TEST(${BACKEND_NAME}, onnx_model_argmin_int32)
     test_case.run();
 }
 
+NGRAPH_TEST(${BACKEND_NAME}, onnx_model_argmax_float)
+{
+    auto function = onnx_import::import_onnx_model(
+        file_util::path_join(SERIALIZED_ZOO, "onnx/argmax_float.prototxt"));
+
+    auto test_case = test::TestCase<TestEngine>(function);
+    test_case.add_input<float>({4, 0.1, 2, 3, -3, 1, -0.9, 0, 1, 2, 3, 0});
+    test_case.add_expected_output<std::int64_t>({0, 3, 0});
+    test_case.run();
+}
+
+NGRAPH_TEST(${BACKEND_NAME}, onnx_model_argmin_float)
+{
+    auto function = onnx_import::import_onnx_model(
+        file_util::path_join(SERIALIZED_ZOO, "onnx/argmin_float.prototxt"));
+
+    auto test_case = test::TestCase<TestEngine>(function);
+    test_case.add_input<float>({4, 0.1, 2, 3, -3, 1, -0.9, 0, 1, 2, 3, 0});
+    test_case.add_expected_output<std::int64_t>({1, 1, 0, 2});
+    test_case.run();
+}
+
 NGRAPH_TEST(${BACKEND_NAME}, onnx_model_top_k)
 {
     auto function =
diff --git a/ngraph/test/runtime/ie/unit_test.manifest b/ngraph/test/runtime/ie/unit_test.manifest
index a3b9f52f82d..266607c43e6 100644
--- a/ngraph/test/runtime/ie/unit_test.manifest
+++ b/ngraph/test/runtime/ie/unit_test.manifest
@@ -154,6 +154,10 @@ onnx_model_argmax_int32
 onnx_model_argmin_int32
 arg_max_dyn_shape
 
+# Result mismatch
+onnx_model_argmax_float
+onnx_model_argmin_float
+
 # Constant has zero dimension that is not allowable
 onnx_dyn_shapes_transpose