Support Half/BFloat16 in masked_fill (#7828)

swolchok · YIWENX14 · commit 492b1b1af846 · 2025-01-28T14:12:15.000-08:00
Partial fix for #7748.
diff --git a/kernels/portable/cpu/op_masked_fill.cpp b/kernels/portable/cpu/op_masked_fill.cpp
@@ -42,8 +42,8 @@ Tensor& masked_fill_scalar_out(
   ET_KERNEL_CHECK(
       ctx, tensors_have_same_dim_order(in, mask, out), InvalidArgument, out);
 
-  ET_SWITCH_REAL_TYPES_AND(
-      Bool, in_type, ctx, "masked_fill.Scalar_out", CTYPE, [&]() {
+  ET_SWITCH_REALHBBF16_TYPES(
+      in_type, ctx, "masked_fill.Scalar_out", CTYPE, [&]() {
         ET_SWITCH_REAL_TYPES_AND(
             Bool, val_type, ctx, "masked_fill.Scalar_out", CTYPE_VAL, [&]() {
               CTYPE_VAL value_v;
diff --git a/kernels/test/op_masked_fill_test.cpp b/kernels/test/op_masked_fill_test.cpp
@@ -114,8 +114,11 @@ TEST_F(OpMaskedFillTest, IntTensorFloatAlphaDies) {
           tf.ones(sizes), tf.ones(sizes), /*alpha=*/.7, out));
 }
 
-TEST_F(OpMaskedFillTest, FloatTensors) {
-  test_floating_point_masked_fill_scalar_out<ScalarType::Float>();
+TEST_F(OpMaskedFillTest, FloatingPointTensors) {
+#define TEST_ENTRY(ctype, dtype) \
+  test_floating_point_masked_fill_scalar_out<ScalarType::dtype>();
+  ET_FORALL_FLOATHBF16_TYPES(TEST_ENTRY);
+#undef TEST_ENTRY
 }
 
 TEST_F(OpMaskedFillTest, DoubleTensors) {

Original file line number	Diff line number	Diff line change
`@@ -114,8 +114,11 @@ TEST_F(OpMaskedFillTest, IntTensorFloatAlphaDies) {`
`114`	`114`	`tf.ones(sizes), tf.ones(sizes), /alpha=/.7, out));`
`115`	`115`	`}`
`116`	`116`
`117`		`-TEST_F(OpMaskedFillTest, FloatTensors) {`
`118`		`- test_floating_point_masked_fill_scalar_out<ScalarType::Float>();`
	`117`	`+TEST_F(OpMaskedFillTest, FloatingPointTensors) {`
	`118`	`+#define TEST_ENTRY(ctype, dtype) \`
	`119`	`+ test_floating_point_masked_fill_scalar_out<ScalarType::dtype>();`
	`120`	`+ ET_FORALL_FLOATHBF16_TYPES(TEST_ENTRY);`
	`121`	`+#undef TEST_ENTRY`
`119`	`122`	`}`
`120`	`123`
`121`	`124`	`TEST_F(OpMaskedFillTest, DoubleTensors) {`