Support Half/BFloat16 in constant_pad_nd (pytorch#7806)

Partial fix for pytorch#7748.

Author: swolchok

kernels/portable/cpu/op_constant_pad_nd.cpp

@@ -184,17 +184,16 @@ Tensor& constant_pad_nd_out(
   ScalarType in_type = in.scalar_type();
   ScalarType value_type = utils::get_scalar_dtype(value);
 
-  ET_SWITCH_REAL_TYPES_AND(
-      Bool, in_type, ctx, "constant_pad_nd.out", CTYPE, [&]() {
-        CTYPE value_v;
-        ET_SWITCH_SCALAR_OBJ_TYPES(
-            value_type, ctx, "constant_pad_nd.out", CTYPE_VALUE, [&]() {
-              CTYPE_VALUE val;
-              utils::extract_scalar(value, &val);
-              value_v = static_cast<CTYPE>(val);
-            });
-        constant_pad_nd_out_impl<CTYPE>(in, pad, value_v, out);
-      });
+  ET_SWITCH_REALHBBF16_TYPES(in_type, ctx, "constant_pad_nd.out", CTYPE, [&]() {
+    CTYPE value_v;
+    ET_SWITCH_SCALAR_OBJ_TYPES(
+        value_type, ctx, "constant_pad_nd.out", CTYPE_VALUE, [&]() {
+          CTYPE_VALUE val;
+          utils::extract_scalar(value, &val);
+          value_v = static_cast<CTYPE>(val);
+        });
+    constant_pad_nd_out_impl<CTYPE>(in, pad, value_v, out);
+  });
 
   return out;
 }

kernels/test/op_constant_pad_nd_test.cpp

@@ -50,7 +50,7 @@ class OpConstantPadNDOutTest : public OperatorTest {
            5,  6,  7,  8,
            1,  2,  3,  4,
            5,  6,  7,  8,
-  
+
            1,  2,  3,  4,
            5,  6,  7,  8,
            1,  2,  3,  4,
@@ -66,7 +66,7 @@ class OpConstantPadNDOutTest : public OperatorTest {
            7,  5,  6,  7,  8,  7,
            7,  1,  2,  3,  4,  7,
            7,  5,  6,  7,  8,  7,
-  
+
            7,  1,  2,  3,  4,  7,
            7,  5,  6,  7,  8,  7,
            7,  1,  2,  3,  4,  7,
@@ -98,7 +98,7 @@ class OpConstantPadNDOutTest : public OperatorTest {
            5,  6,  7,  8,
            1,  2,  3,  4,
            5,  6,  7,  8,
-  
+
            1,  2,  3,  4,
            5,  6,  7,  8,
            1,  2,  3,  4,
@@ -116,7 +116,7 @@ class OpConstantPadNDOutTest : public OperatorTest {
            5,  6,  7,  8,
            1,  2,  3,  4,
            5,  6,  7,  8,
-  
+
            7,  7,  7,  7,
            7,  7,  7,  7,
            1,  2,  3,  4,
@@ -150,7 +150,7 @@ class OpConstantPadNDOutTest : public OperatorTest {
            5,  6,  7,  8,
            1,  2,  3,  4,
            5,  6,  7,  8,
-  
+
            1,  2,  3,  4,
            5,  6,  7,  8,
            1,  2,  3,  4,
@@ -166,12 +166,12 @@ class OpConstantPadNDOutTest : public OperatorTest {
            7,  7,  7,  7,
            7,  7,  7,  7,
            7,  7,  7,  7,
-  
+
            1,  2,  3,  4,
            5,  6,  7,  8,
            1,  2,  3,  4,
            5,  6,  7,  8,
-  
+
            1,  2,  3,  4,
            5,  6,  7,  8,
            1,  2,  3,  4,
@@ -203,7 +203,7 @@ class OpConstantPadNDOutTest : public OperatorTest {
            5,  6,  7,  8,
            1,  2,  3,  4,
            5,  6,  7,  8,
-  
+
            1,  2,  3,  4,
            5,  6,  7,  8,
            1,  2,  3,  4,
@@ -221,7 +221,7 @@ class OpConstantPadNDOutTest : public OperatorTest {
            7,  7,  5,  6,  7,  8,  7,
            7,  7,  7,  7,  7,  7,  7,
            7,  7,  7,  7,  7,  7,  7,
-  
+
            7,  7,  1,  2,  3,  4,  7,
            7,  7,  5,  6,  7,  8,  7,
            7,  7,  1,  2,  3,  4,  7,
@@ -255,7 +255,7 @@ class OpConstantPadNDOutTest : public OperatorTest {
            5,  6,  7,  8,
            1,  2,  3,  4,
            5,  6,  7,  8,
-  
+
            1,  2,  3,  4,
            5,  6,  7,  8,
            1,  2,  3,  4,
@@ -271,12 +271,12 @@ class OpConstantPadNDOutTest : public OperatorTest {
            7,  7,  5,  6,  7,  8,  7,
            7,  7,  1,  2,  3,  4,  7,
            7,  7,  5,  6,  7,  8,  7,
-  
+
            7,  7,  1,  2,  3,  4,  7,
            7,  7,  5,  6,  7,  8,  7,
            7,  7,  1,  2,  3,  4,  7,
            7,  7,  5,  6,  7,  8,  7,
-  
+
            7,  7,  7,  7,  7,  7,  7,
            7,  7,  7,  7,  7,  7,  7,
            7,  7,  7,  7,  7,  7,  7,
@@ -308,7 +308,7 @@ class OpConstantPadNDOutTest : public OperatorTest {
            5,  6,  7,  8,
            1,  2,  3,  4,
            5,  6,  7,  8,
-  
+
            1,  2,  3,  4,
            5,  6,  7,  8,
            1,  2,  3,  4,
@@ -325,13 +325,13 @@ class OpConstantPadNDOutTest : public OperatorTest {
            7,  7,  5,  6,  7,  8,  7,
            7,  7,  1,  2,  3,  4,  7,
            7,  7,  5,  6,  7,  8,  7,
-  
+
            7,  7,  7,  7,  7,  7,  7,
            7,  7,  1,  2,  3,  4,  7,
            7,  7,  5,  6,  7,  8,  7,
            7,  7,  1,  2,  3,  4,  7,
            7,  7,  5,  6,  7,  8,  7,
-  
+
            7,  7,  7,  7,  7,  7,  7,
            7,  7,  7,  7,  7,  7,  7,
            7,  7,  7,  7,  7,  7,  7,
@@ -353,47 +353,47 @@ TEST_F(OpConstantPadNDOutTest, TestPadDim2) {
 #define TEST_ENTRY(ctype, dtype) \
   test_constant_pad_nd_out_dim2<ScalarType::dtype>();
 
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }
 
 TEST_F(OpConstantPadNDOutTest, TestPadDim1) {
 #define TEST_ENTRY(ctype, dtype) \
   test_constant_pad_nd_out_dim1<ScalarType::dtype>();
 
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }
 
 TEST_F(OpConstantPadNDOutTest, TestPadDim0) {
 #define TEST_ENTRY(ctype, dtype) \
   test_constant_pad_nd_out_dim0<ScalarType::dtype>();
 
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }
 
 TEST_F(OpConstantPadNDOutTest, TestPadDim1And2) {
 #define TEST_ENTRY(ctype, dtype) \
   test_constant_pad_nd_out_dim12<ScalarType::dtype>();
 
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }
 
 TEST_F(OpConstantPadNDOutTest, TestPadDim0And2) {
 #define TEST_ENTRY(ctype, dtype) \
   test_constant_pad_nd_out_dim02<ScalarType::dtype>();
 
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }
 
 TEST_F(OpConstantPadNDOutTest, TestPadDim0And1And2) {
 #define TEST_ENTRY(ctype, dtype) \
   test_constant_pad_nd_out_dim012<ScalarType::dtype>();
 
-  ET_FORALL_REAL_TYPES(TEST_ENTRY);
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }