Support Half/BFloat16 in scatter

Partial fix for #7748.

ghstack-source-id: 40b6cdd
ghstack-comment-id: 2608589452
Pull Request resolved: #7864

Author: swolchok

kernels/portable/cpu/op_scatter.cpp

@@ -125,7 +125,7 @@ Tensor& scatter_src_out(
 
   constexpr auto name = "scatter.src_out";
 
-  ET_SWITCH_REALHB_TYPES(in.scalar_type(), ctx, name, CTYPE, [&]() {
+  ET_SWITCH_REALHBBF16_TYPES(in.scalar_type(), ctx, name, CTYPE, [&]() {
     scatter_src_helper<CTYPE>(in, dim, index, src, out);
   });
 
@@ -158,7 +158,7 @@ Tensor& scatter_value_out(
     CTYPE_VAL val;
     utils::extract_scalar(value, &val);
 
-    ET_SWITCH_REALHB_TYPES(in.scalar_type(), ctx, name, CTYPE, [&]() {
+    ET_SWITCH_REALHBBF16_TYPES(in.scalar_type(), ctx, name, CTYPE, [&]() {
       scatter_value_helper<CTYPE>(in, dim, index, val, out);
     });
   });

kernels/test/CMakeLists.txt

@@ -201,6 +201,7 @@ set(all_test_sources
     "op_rsqrt_test.cpp"
     "op_rsub_test.cpp"
     "op_scalar_tensor_test.cpp"
+    "op_scatter_test.cpp"
     "op_scatter_add_test.cpp"
     "op_select_scatter_test.cpp"
     "op_select_copy_test.cpp"

kernels/test/op_scatter_test.cpp

@@ -368,7 +368,7 @@ class OpScatterValueOutTest : public OperatorTest {
 
 TEST_F(OpScatterSrcOutTest, AllValidInputOutputSupport) {
 #define TEST_ENTRY(CTYPE, DTYPE) test_scatter_src_out<ScalarType::DTYPE>();
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }
 
@@ -381,7 +381,7 @@ TEST_F(OpScatterSrcOutTest, InvalidDimensionsDies) {
 
 TEST_F(OpScatterValueOutTest, AllValidInputOutputSupport) {
 #define TEST_ENTRY(CTYPE, DTYPE) test_scatter_value_out<ScalarType::DTYPE>();
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }