Support BFloat16 in gather

Partial fix for #7748.

ghstack-source-id: 8d6b844
ghstack-comment-id: 2605979600
Pull Request resolved: #7823

Author: swolchok

kernels/portable/cpu/op_gather.cpp

@@ -86,7 +86,7 @@ Tensor& gather_out(
 
   constexpr auto name = "gather.out";
 
-  ET_SWITCH_REALHB_TYPES(in.scalar_type(), ctx, name, CTYPE, [&]() {
+  ET_SWITCH_REALHBBF16_TYPES(in.scalar_type(), ctx, name, CTYPE, [&]() {
     gather_helper<CTYPE>(in, index, out, dim);
   });
 

kernels/test/CMakeLists.txt

@@ -139,6 +139,7 @@ set(all_test_sources
     "op_fmod_test.cpp"
     "op_full_like_test.cpp"
     "op_full_test.cpp"
+    "op_gather_test.cpp"
     "op_ge_test.cpp"
     "op_gelu_test.cpp"
     "op_glu_test.cpp"

kernels/test/op_gather_test.cpp

@@ -194,7 +194,7 @@ class OpGatherOutTest : public OperatorTest {
 
 TEST_F(OpGatherOutTest, AllValidInputOutputSupport) {
 #define TEST_ENTRY(CTYPE, DTYPE) test_gather_out<ScalarType::DTYPE>();
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }