Support Half/BFloat16 in stack

Partial fix for #7748.

ghstack-source-id: 4a0c4e2
ghstack-comment-id: 2610826380
Pull Request resolved: #7894

Author: swolchok

kernels/portable/cpu/op_stack.cpp

@@ -55,21 +55,20 @@ Tensor& stack_out(
   const size_t ninputs = tensors.size();
 
   const auto out_type = out.scalar_type();
-  ET_SWITCH_REAL_TYPES_AND(Bool, out_type, ctx, "stack.out", CTYPE_OUT, [&] {
+  ET_SWITCH_REALHBBF16_TYPES(out_type, ctx, "stack.out", CTYPE_OUT, [&] {
     CTYPE_OUT* out_ptr = out.mutable_data_ptr<CTYPE_OUT>();
     for (size_t i = 0; i < outer; ++i) {
       for (size_t j = 0; j < ninputs; ++j) {
         const auto in_type = tensors[j].scalar_type();
-        ET_SWITCH_REAL_TYPES_AND(
-            Bool, in_type, ctx, "stack.out", CTYPE_IN, [&] {
-              const CTYPE_IN* const in_ptr =
-                  tensors[j].const_data_ptr<CTYPE_IN>() + i * inner;
+        ET_SWITCH_REALHBBF16_TYPES(in_type, ctx, "stack.out", CTYPE_IN, [&] {
+          const CTYPE_IN* const in_ptr =
+              tensors[j].const_data_ptr<CTYPE_IN>() + i * inner;
 
-              for (size_t k = 0; k < inner; ++k) {
-                out_ptr[k] = static_cast<CTYPE_OUT>(in_ptr[k]);
-              }
-              out_ptr += inner;
-            });
+          for (size_t k = 0; k < inner; ++k) {
+            out_ptr[k] = static_cast<CTYPE_OUT>(in_ptr[k]);
+          }
+          out_ptr += inner;
+        });
       }
     }
   });

kernels/test/op_stack_test.cpp

@@ -276,7 +276,7 @@ TEST_F(OpStackOutTest, InsertEnd) {
 /// zeros().
 TEST_F(OpStackOutTest, AllDtypesSupported) {
 #define TEST_ENTRY(ctype, dtype) test_dtype<ctype, ScalarType::dtype>();
-  ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
+  ET_FORALL_REALHBBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
   // TODO: Also add tests for half, complex, quantized, and other types. Easiest
   // way to do that would be to make TensorFactory support zeros() and ones()