Update

[ghstack-poisoned]

Author: swolchok

kernels/portable/cpu/op_mean.cpp

@@ -44,8 +44,8 @@ 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, [&] {
+  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) {

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
 }