Support Half/BFloat16 in arange (#7791)

Partial fix for #7748.

Author: swolchok

kernels/portable/cpu/op_arange.cpp

@@ -39,7 +39,7 @@ Tensor& arange_out(KernelRuntimeContext& ctx, const Scalar& end, Tensor& out) {
       InvalidArgument,
       out);
 
-  ET_SWITCH_REAL_TYPES(out.scalar_type(), ctx, "arange.out", CTYPE, [&]() {
+  ET_SWITCH_REALHBF16_TYPES(out.scalar_type(), ctx, "arange.out", CTYPE, [&]() {
     auto out_data = out.mutable_data_ptr<CTYPE>();
     for (size_t i = 0; i < size; i++) {
       out_data[i] = static_cast<CTYPE>(i);
@@ -88,7 +88,7 @@ Tensor& arange_start_out(
       InvalidArgument,
       out);
 
-  ET_SWITCH_REAL_TYPES(
+  ET_SWITCH_REALHBF16_TYPES(
       out.scalar_type(), ctx, "arange.start_out", CTYPE, [&]() {
         auto out_data = out.mutable_data_ptr<CTYPE>();
         for (size_t i = 0; i < size; i++) {

kernels/test/op_arange_test.cpp

@@ -86,9 +86,9 @@ class OpArangeStartOutTest : public OperatorTest {
 };
 
 /// A generic smoke test that works for any dtype that supports  zeros().
-TEST_F(OpArangeOutTest, AllRealDtypesSupported) {
+TEST_F(OpArangeOutTest, AllRealHBF16DtypesSupported) {
 #define TEST_ENTRY(ctype, dtype) test_arange_dtype<ctype, ScalarType::dtype>();
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }
 
@@ -164,10 +164,10 @@ TEST_F(OpArangeOutTest, DynamicShapeUnbound) {
 }
 
 /// A generic smoke test that works for any dtype that supports  zeros().
-TEST_F(OpArangeStartOutTest, AllRealDtypesSupported) {
+TEST_F(OpArangeStartOutTest, AllRealHBF16DtypesSupported) {
 #define TEST_ENTRY(ctype, dtype) \
   test_arange_start_dtype<ctype, ScalarType::dtype>();
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }