fix reduce_arg's test (#1046)

* fix reduce_arg's test * fix reduce_arg --------- Co-authored-by: hejunchao <hejunchao@canaan-creative.com>
kendryte · Aug 11, 2023 · ab48a3f · ab48a3f
1 parent 1747ae7
commit ab48a3f
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 20 deletions.
diff --git a/src/Native/src/kernels/stackvm/reference/reduce_arg.cpp b/src/Native/src/kernels/stackvm/reference/reduce_arg.cpp
@@ -60,10 +60,8 @@ result<void> reduce_arg_impl(TReducer &&reducer, T init_value, const T *input,
             out_map[out_idx].clear();
             out_map[out_idx].push_back(index[axes[0]]);
             dst = src;
-        } else if constexpr (std::is_same_v<T, float>) {
-            if (fabs(src - dst) < epsilon) {
-                out_map[out_idx].push_back(index[axes[0]]);
-            }
+        } else if (std::fabs(src - dst) < epsilon) {
+            out_map[out_idx].push_back(index[axes[0]]);
         }
         return ok();
     }));

diff --git a/tests/kernels/test_reduce_arg_max.cpp b/tests/kernels/test_reduce_arg_max.cpp
@@ -78,13 +78,13 @@ class ReduceArgMaxTest : public KernelTest,
 
 INSTANTIATE_TEST_SUITE_P(
     ReduceArgMax, ReduceArgMaxTest,
-    testing::Combine(testing::Values(dt_float32), testing::Values(dt_int64),
-                     testing::Values(dims_t{1, 3, 16, 16}, dims_t{1, 2, 3, 4},
-                                     dims_t{1, 3, 16}, dims_t{3, 16},
-                                     dims_t{16}),
-                     testing::Values(dims_t{1}),
-                     testing::Values(-1, 0, 1, 2, 3, -2, -3, -4),
-                     testing::Values(1, 0), testing::Values(1, 0)));
+    testing::Combine(
+        testing::Values(dt_float64, dt_float32, dt_int32 /*, dt_int64*/),
+        testing::Values(dt_int64),
+        testing::Values(dims_t{1, 3, 16, 16}, dims_t{1, 2, 3, 4},
+                        dims_t{1, 3, 16}, dims_t{3, 16}, dims_t{16}),
+        testing::Values(dims_t{1}), testing::Values(-1, 0, 1, 2, 3, -2, -3, -4),
+        testing::Values(1, 0), testing::Values(1, 0)));
 
 TEST_P(ReduceArgMaxTest, ReduceArgMax) {
 
@@ -95,7 +95,7 @@ TEST_P(ReduceArgMaxTest, ReduceArgMax) {
     void *ptr_ort = tensor_buffer(output_ort, &size);
     dims_t shape(tensor_rank(output_ort));
     tensor_shape(output_ort, reinterpret_cast<int64_t *>(shape.data()));
-    auto expected = hrt::create(dt_float64, shape,
+    auto expected = hrt::create(dt_int64, shape,
                                 {reinterpret_cast<gsl::byte *>(ptr_ort), size},
                                 true, host_runtime_tensor::pool_cpu_only)
                         .expect("create tensor failed");
@@ -112,6 +112,8 @@ TEST_P(ReduceArgMaxTest, ReduceArgMax) {
                   cosine_similarity_tensor(expected, actual);
 
     if (!result) {
+        std::cout << "input ";
+        print_runtime_tensor(a);
         std::cout << "actual ";
         print_runtime_tensor(actual);
         std::cout << "expected ";

diff --git a/tests/kernels/test_reduce_arg_min.cpp b/tests/kernels/test_reduce_arg_min.cpp
@@ -78,13 +78,13 @@ class ReduceArgMinTest : public KernelTest,
 
 INSTANTIATE_TEST_SUITE_P(
     ReduceArgMin, ReduceArgMinTest,
-    testing::Combine(testing::Values(dt_float32), testing::Values(dt_int64),
-                     testing::Values(dims_t{1, 3, 16, 16}, dims_t{1, 2, 3, 4},
-                                     dims_t{1, 3, 16}, dims_t{3, 16},
-                                     dims_t{16}),
-                     testing::Values(dims_t{1}),
-                     testing::Values(-1, 0, 1, 2, 3, -2, -3, -4),
-                     testing::Values(1, 0), testing::Values(1, 0)));
+    testing::Combine(
+        testing::Values(dt_float64, dt_float32, dt_int32 /*, dt_int64*/),
+        testing::Values(dt_int64),
+        testing::Values(dims_t{1, 3, 16, 16}, dims_t{1, 2, 3, 4},
+                        dims_t{1, 3, 16}, dims_t{3, 16}, dims_t{16}),
+        testing::Values(dims_t{1}), testing::Values(-1, 0, 1, 2, 3, -2, -3, -4),
+        testing::Values(1, 0), testing::Values(1, 0)));
 
 TEST_P(ReduceArgMinTest, ReduceArgMin) {
 
@@ -95,7 +95,7 @@ TEST_P(ReduceArgMinTest, ReduceArgMin) {
     void *ptr_ort = tensor_buffer(output_ort, &size);
     dims_t shape(tensor_rank(output_ort));
     tensor_shape(output_ort, reinterpret_cast<int64_t *>(shape.data()));
-    auto expected = hrt::create(dt_float64, shape,
+    auto expected = hrt::create(dt_int64, shape,
                                 {reinterpret_cast<gsl::byte *>(ptr_ort), size},
                                 true, host_runtime_tensor::pool_cpu_only)
                         .expect("create tensor failed");
@@ -112,6 +112,8 @@ TEST_P(ReduceArgMinTest, ReduceArgMin) {
                   cosine_similarity_tensor(expected, actual);
 
     if (!result) {
+        std::cout << "input ";
+        print_runtime_tensor(a);
         std::cout << "actual ";
         print_runtime_tensor(actual);
         std::cout << "expected ";