Support BFloat16 in full

swolchok · swolchok · commit 3b3647086bdd · 2025-01-21T15:47:02.000-08:00
Partial fix for #7748. ghstack-source-id: 55ca1ba ghstack-comment-id: 2605970943 Pull Request resolved: #7821
diff --git a/kernels/portable/cpu/op_full.cpp b/kernels/portable/cpu/op_full.cpp
@@ -40,7 +40,7 @@ Tensor& full_out(
     CTYPE_VAL val;
     utils::extract_scalar(fill_value, &val);
 
-    ET_SWITCH_REALHB_TYPES(out_type, ctx, name, CTYPE_OUT, [&] {
+    ET_SWITCH_REALHBBF16_TYPES(out_type, ctx, name, CTYPE_OUT, [&] {
       CTYPE_OUT val_casted = static_cast<CTYPE_OUT>(val);
       auto data_out = out.mutable_data_ptr<CTYPE_OUT>();
       for (size_t i = 0; i < out.numel(); ++i) {
diff --git a/kernels/test/op_full_test.cpp b/kernels/test/op_full_test.cpp
@@ -70,7 +70,7 @@ class OpFullOutTest : public OperatorTest {
     test_ones_out<ScalarType::DTYPE>({2, 3, 4}); \
   }
 
-ET_FORALL_REALH_TYPES(GENERATE_TEST)
+ET_FORALL_REALHBF16_TYPES(GENERATE_TEST)
 
 TEST_F(OpFullOutTest, ValueOverflow) {
   if (torch::executor::testing::SupportedFeatures::get()->is_aten) {

Original file line number	Diff line number	Diff line change
`@@ -70,7 +70,7 @@ class OpFullOutTest : public OperatorTest {`
`70`	`70`	`test_ones_out<ScalarType::DTYPE>({2, 3, 4}); \`
`71`	`71`	`}`
`72`	`72`
`73`		`-ET_FORALL_REALH_TYPES(GENERATE_TEST)`
	`73`	`+ET_FORALL_REALHBF16_TYPES(GENERATE_TEST)`
`74`	`74`
`75`	`75`	`TEST_F(OpFullOutTest, ValueOverflow) {`
`76`	`76`	`if (torch::executor::testing::SupportedFeatures::get()->is_aten) {`