Support Half/BFloat16 in ones (#7851)

swolchok · YIWENX14 · commit a561e31e11c8 · 2025-01-28T14:20:31.000-08:00
Partial fix for #7748.
diff --git a/kernels/portable/cpu/op_ones.cpp b/kernels/portable/cpu/op_ones.cpp
@@ -20,7 +20,7 @@ Tensor& ones_out(KernelRuntimeContext& ctx, IntArrayRef size, Tensor& out) {
       ctx, resize_tensor(out, size) == Error::Ok, InvalidArgument, out);
 
   ScalarType out_type = out.scalar_type();
-  ET_SWITCH_REAL_TYPES_AND(Bool, out_type, ctx, __func__, CTYPE, [&] {
+  ET_SWITCH_REALHBBF16_TYPES(out_type, ctx, __func__, CTYPE, [&] {
     auto out_data = out.mutable_data_ptr<CTYPE>();
     for (size_t i = 0; i < out.numel(); i++) {
       out_data[i] = static_cast<CTYPE>(1);
diff --git a/kernels/test/op_ones_test.cpp b/kernels/test/op_ones_test.cpp
@@ -51,4 +51,4 @@ class OpOnesOutTest : public OperatorTest {
     test_ones_out<ScalarType::DTYPE>({2, 3, 4}); \
   }
 
-ET_FORALL_REAL_TYPES_AND(Bool, GENERATE_TEST)
+ET_FORALL_REALHBBF16_TYPES(GENERATE_TEST)

Original file line number	Diff line number	Diff line change
`@@ -51,4 +51,4 @@ class OpOnesOutTest : public OperatorTest {`
`51`	`51`	`test_ones_out<ScalarType::DTYPE>({2, 3, 4}); \`
`52`	`52`	`}`
`53`	`53`
`54`		`-ET_FORALL_REAL_TYPES_AND(Bool, GENERATE_TEST)`
	`54`	`+ET_FORALL_REALHBBF16_TYPES(GENERATE_TEST)`