Update

[ghstack-poisoned]

Author: swolchok

kernels/portable/cpu/op_argmax.cpp

@@ -43,7 +43,7 @@ Tensor& argmax_out(
   ET_KERNEL_CHECK(
       ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
 
-  ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, "argmax.out", CTYPE, [&] {
+  ET_SWITCH_REALHBF16_TYPES(in.scalar_type(), ctx, "argmax.out", CTYPE, [&] {
     long* out_data = out.mutable_data_ptr<long>();
 
     for (size_t out_ix = 0; out_ix < out.numel(); ++out_ix) {

kernels/portable/cpu/op_argmin.cpp

@@ -43,7 +43,7 @@ Tensor& argmin_out(
   ET_KERNEL_CHECK(
       ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
 
-  ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, "argmin.out", CTYPE, [&] {
+  ET_SWITCH_REALHBF16_TYPES(in.scalar_type(), ctx, "argmin.out", CTYPE, [&] {
     long* out_data = out.mutable_data_ptr<long>();
 
     for (size_t out_ix = 0; out_ix < out.numel(); ++out_ix) {

kernels/test/op_argmax_test.cpp

@@ -31,57 +31,39 @@ class OpArgmaxTest : public OperatorTest {
       Tensor& out) {
     return torch::executor::aten::argmax_outf(context_, in, dim, keepdim, out);
   }
-};
-
-TEST_F(OpArgmaxTest, SanityCheckLong) {
-  TensorFactory<ScalarType::Long> tf;
-
-  // clang-format off
-  Tensor in = tf.make(
-    { 2, 3, 4 },
-    { 1, 4, 1, 6,
-      5, 8, 5, 6,
-      5, 3, 9, 2,
-
-      3, 9, 1, 4,
-      9, 7, 5, 5,
-      7, 7, 6, 3 });
-
-  Tensor out = tf.zeros({2, 4});
-  Tensor expected = tf.make({2, 4}, {
-    1, 1, 2, 0,
-    1, 0, 2, 1 });
-  Tensor ret = op_argmax_out(in, 1, false, out);
-
-  EXPECT_TENSOR_EQ(out, ret);
-  EXPECT_TENSOR_EQ(out, expected);
-  // clang-format on
-}
-
-TEST_F(OpArgmaxTest, SanityCheckShort) {
-  TensorFactory<ScalarType::Long> tfl;
-  TensorFactory<ScalarType::Short> tfs;
 
-  // clang-format off
-  Tensor in = tfs.make(
-    { 2, 3, 4 },
-    { 1, 4, 1, 6,
-      5, 8, 5, 6,
-      5, 3, 9, 2,
-
-      3, 9, 1, 4,
-      9, 7, 5, 5,
-      7, 7, 6, 3 });
-
-  Tensor out = tfl.zeros({2, 4});
-  Tensor expected = tfl.make({2, 4}, {
-    1, 1, 2, 0,
-    1, 0, 2, 1 });
-  Tensor ret = op_argmax_out(in, 1, false, out);
+  template <ScalarType dtype>
+  void test_argmax_dtype() {
+    TensorFactory<ScalarType::Long> tfl;
+    TensorFactory<dtype> tf_dtype;
+
+    // clang-format off
+    Tensor in = tf_dtype.make(
+        { 2, 3, 4 },
+        { 1, 4, 1, 6,
+          5, 8, 5, 6,
+          5, 3, 9, 2,
+
+          3, 9, 1, 4,
+          9, 7, 5, 5,
+          7, 7, 6, 3 });
+
+    Tensor out = tfl.zeros({2, 4});
+    Tensor expected = tfl.make({2, 4}, {
+        1, 1, 2, 0,
+        1, 0, 2, 1 });
+    Tensor ret = op_argmax_out(in, 1, false, out);
+
+    EXPECT_TENSOR_EQ(out, ret);
+    EXPECT_TENSOR_EQ(out, expected);
+    // clang-format on
+  }
+};
 
-  EXPECT_TENSOR_EQ(out, ret);
-  EXPECT_TENSOR_EQ(out, expected);
-  // clang-format on
+TEST_F(OpArgmaxTest, SanityCheck) {
+#define TEST_ENTRY(ctype, dtype) test_argmax_dtype<ScalarType::dtype>();
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
+#undef TEST_ENTRY
 }
 
 TEST_F(OpArgmaxTest, SanityCheckNullDim) {

kernels/test/op_argmin_test.cpp

@@ -31,57 +31,39 @@ class OpArgminTest : public OperatorTest {
       Tensor& out) {
     return torch::executor::aten::argmin_outf(context_, in, dim, keepdim, out);
   }
-};
-
-TEST_F(OpArgminTest, SanityCheckLong) {
-  TensorFactory<ScalarType::Long> tf;
-
-  // clang-format off
-  Tensor in = tf.make(
-    { 2, 3, 4 },
-    { 1, 4, 1, 6,
-      5, 8, 5, 6,
-      5, 3, 9, 2,
-
-      3, 9, 1, 4,
-      9, 7, 5, 5,
-      7, 7, 6, 3 });
-
-  Tensor out = tf.zeros({2, 4});
-  Tensor expected = tf.make({2, 4}, {
-    0, 2, 0, 2,
-    0, 1, 0, 2 });
-  Tensor ret = op_argmin_out(in, 1, false, out);
-
-  EXPECT_TENSOR_EQ(out, ret);
-  EXPECT_TENSOR_EQ(out, expected);
-  // clang-format on
-}
-
-TEST_F(OpArgminTest, SanityCheckShort) {
-  TensorFactory<ScalarType::Long> tfl;
-  TensorFactory<ScalarType::Short> tfs;
 
-  // clang-format off
-  Tensor in = tfs.make(
-    { 2, 3, 4 },
-    { 1, 4, 1, 6,
-      5, 8, 5, 6,
-      5, 3, 9, 2,
-
-      3, 9, 1, 4,
-      9, 7, 5, 5,
-      7, 7, 6, 3 });
-
-  Tensor out = tfl.zeros({2, 4});
-  Tensor expected = tfl.make({2, 4}, {
-    0, 2, 0, 2,
-    0, 1, 0, 2 });
-  Tensor ret = op_argmin_out(in, 1, false, out);
+  template <ScalarType dtype>
+  void test_argmin_dtype() {
+    TensorFactory<ScalarType::Long> tfl;
+    TensorFactory<dtype> tf_dtype;
+
+    // clang-format off
+    Tensor in = tf_dtype.make(
+        { 2, 3, 4 },
+        { 1, 4, 1, 6,
+          5, 8, 5, 6,
+          5, 3, 9, 2,
+
+          3, 9, 1, 4,
+          9, 7, 5, 5,
+          7, 7, 6, 3 });
+
+    Tensor out = tfl.zeros({2, 4});
+    Tensor expected = tfl.make({2, 4}, {
+        0, 2, 0, 2,
+        0, 1, 0, 2 });
+    Tensor ret = op_argmin_out(in, 1, false, out);
+
+    EXPECT_TENSOR_EQ(out, ret);
+    EXPECT_TENSOR_EQ(out, expected);
+    // clang-format on
+  }
+};
 
-  EXPECT_TENSOR_EQ(out, ret);
-  EXPECT_TENSOR_EQ(out, expected);
-  // clang-format on
+TEST_F(OpArgminTest, SanityCheck) {
+#define TEST_ENTRY(ctype, dtype) test_argmin_dtype<ScalarType::dtype>();
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
+#undef TEST_ENTRY
 }
 
 TEST_F(OpArgminTest, SanityCheckNullDim) {