Support bfloat16 in op_index (pytorch#5499)

Summary:
Pull Request resolved: pytorch#5499

Seems to block bfloat16 stories110M as exported by torchchat (and we should have op coverage for bfloat16 anyway).
ghstack-source-id: 243857968

Reviewed By: larryliu0820

Differential Revision: D63054001

fbshipit-source-id: 530b479872643f878912592c7b260d71e6e05804

Author: swolchok

kernels/portable/cpu/op_index.cpp

@@ -89,7 +89,7 @@ Tensor& index_Tensor_out(
   compute_dim_map(in, indices, dim_map, block_count == 1);
   compute_index_map(in, indices, ix_map);
 
-  ET_SWITCH_REALHB_TYPES(in_type, ctx, "index.Tensor_out", CTYPE, [&]() {
+  ET_SWITCH_REALHBBF16_TYPES(in_type, ctx, "index.Tensor_out", CTYPE, [&]() {
     const CTYPE* const in_data = in.const_data_ptr<CTYPE>();
     CTYPE* const out_data = out.mutable_data_ptr<CTYPE>();
 

kernels/test/op_index_test.cpp

@@ -107,7 +107,7 @@ class OpIndexTensorOutTest : public OperatorTest {
 #define TEST_ENTRY(ctype, dtype) \
   test_dtype<ScalarType::dtype, ScalarType::Long, ScalarType::dtype>();
 
-    ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
+    ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 
 #undef TEST_ENTRY
   }