Support Half/BFloat16 in mean (#7837)

Partial fix for #7748.

Author: swolchok

kernels/portable/cpu/op_mean.cpp

@@ -44,23 +44,24 @@ Tensor& mean_dim_out(
       InvalidArgument,
       out);
 
-  ET_SWITCH_REALHB_TYPES(in.scalar_type(), ctx, "mean.out", CTYPE_IN, [&] {
-    ET_SWITCH_FLOATH_TYPES(out.scalar_type(), ctx, "mean.out", CTYPE_OUT, [&] {
-      CTYPE_OUT* out_data = out.mutable_data_ptr<CTYPE_OUT>();
-      const size_t num = get_reduced_dim_product(in, dim_list);
-      for (size_t out_ix = 0; out_ix < out.numel(); ++out_ix) {
-        CTYPE_OUT sum = 0;
-        if (in.numel() > 0) {
-          sum = map_reduce_over_dim_list<CTYPE_IN, CTYPE_OUT>(
-              [](CTYPE_IN v) { return static_cast<CTYPE_OUT>(v); },
-              [](CTYPE_OUT outv, CTYPE_OUT acc) { return acc + outv; },
-              in,
-              dim_list,
-              out_ix);
-        }
-        out_data[out_ix] = sum / static_cast<float>(num);
-      }
-    });
+  ET_SWITCH_REALHBBF16_TYPES(in.scalar_type(), ctx, "mean.out", CTYPE_IN, [&] {
+    ET_SWITCH_FLOATHBF16_TYPES(
+        out.scalar_type(), ctx, "mean.out", CTYPE_OUT, [&] {
+          CTYPE_OUT* out_data = out.mutable_data_ptr<CTYPE_OUT>();
+          const size_t num = get_reduced_dim_product(in, dim_list);
+          for (size_t out_ix = 0; out_ix < out.numel(); ++out_ix) {
+            CTYPE_OUT sum = 0;
+            if (in.numel() > 0) {
+              sum = map_reduce_over_dim_list<CTYPE_IN, CTYPE_OUT>(
+                  [](CTYPE_IN v) { return static_cast<CTYPE_OUT>(v); },
+                  [](CTYPE_OUT outv, CTYPE_OUT acc) { return acc + outv; },
+                  in,
+                  dim_list,
+                  out_ix);
+            }
+            out_data[out_ix] = sum / static_cast<float>(num);
+          }
+        });
   });
 
   return out;

kernels/test/op_mean_test.cpp

@@ -238,6 +238,18 @@ class OpMeanOutTest : public OperatorTest {
   }
 };
 
+template <>
+void OpMeanOutTest::
+    test_mean_dim_out_dtype<ScalarType::Bool, ScalarType::Half>() {
+  test_mean_dim_out_bool<ScalarType::Half>();
+}
+
+template <>
+void OpMeanOutTest::
+    test_mean_dim_out_dtype<ScalarType::Bool, ScalarType::BFloat16>() {
+  test_mean_dim_out_bool<ScalarType::BFloat16>();
+}
+
 template <>
 void OpMeanOutTest::
     test_mean_dim_out_dtype<ScalarType::Bool, ScalarType::Float>() {
@@ -331,9 +343,9 @@ TEST_F(OpMeanOutTest, AllRealInputFloatOutputPasses) {
   test_mean_dim_out_dtype<ScalarType::INPUT_DTYPE, ScalarType::OUTPUT_DTYPE>();
 
 #define TEST_ENTRY(INPUT_CTYPE, INPUT_DTYPE) \
-  ET_FORALL_FLOAT_TYPES_WITH2(INPUT_CTYPE, INPUT_DTYPE, TEST_KERNEL);
+  ET_FORALL_FLOATHBF16_TYPES_WITH2(INPUT_CTYPE, INPUT_DTYPE, TEST_KERNEL);
 
-  ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
+  ET_FORALL_REALHBBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 #undef TEST_KERNEL
 }