Add Half support: full.out

Differential Revision: D61864622

Pull Request resolved: #4934

Author: manuelcandales

kernels/portable/cpu/op_full.cpp

@@ -34,11 +34,13 @@ Tensor& full_out(
       out,
       "Failed to resize output tensor.");
 
-  ET_SWITCH_REAL_TYPES_AND(Bool, val_type, ctx, "full.out", CTYPE_VAL, [&] {
+  constexpr auto name = "full.out";
+
+  ET_SWITCH_SCALAR_OBJ_TYPES(val_type, ctx, name, CTYPE_VAL, [&] {
     CTYPE_VAL val;
     utils::extract_scalar(fill_value, &val);
 
-    ET_SWITCH_REAL_TYPES_AND(Bool, out_type, ctx, "full.out", CTYPE_OUT, [&] {
+    ET_SWITCH_REALHB_TYPES(out_type, ctx, name, CTYPE_OUT, [&] {
       CTYPE_OUT val_casted = static_cast<CTYPE_OUT>(val);
       auto data_out = out.mutable_data_ptr<CTYPE_OUT>();
       for (size_t i = 0; i < out.numel(); ++i) {

kernels/test/op_full_test.cpp

@@ -38,12 +38,25 @@ class OpFullOutTest : public OperatorTest {
     std::vector<int64_t> size_int64_t(size_int32_t.begin(), size_int32_t.end());
     auto aref = IntArrayRef(size_int64_t.data(), size_int64_t.size());
 
+    // Boolean Scalar
     // Before: `out` consists of 0s.
     Tensor out = tf.zeros(size_int32_t);
+    // After: `out` consists of 1s.
+    op_full_out(aref, true, out);
+    EXPECT_TENSOR_EQ(out, tf.ones(size_int32_t));
 
+    // Integral Scalar
+    // Before: `out` consists of 0s.
+    out = tf.zeros(size_int32_t);
     // After: `out` consists of 1s.
     op_full_out(aref, 1, out);
+    EXPECT_TENSOR_EQ(out, tf.ones(size_int32_t));
 
+    // Floating Point Scalar
+    // Before: `out` consists of 0s.
+    out = tf.zeros(size_int32_t);
+    // After: `out` consists of 1s.
+    op_full_out(aref, 1.0, out);
     EXPECT_TENSOR_EQ(out, tf.ones(size_int32_t));
   }
 };
@@ -57,4 +70,55 @@ class OpFullOutTest : public OperatorTest {
     test_ones_out<ScalarType::DTYPE>({2, 3, 4}); \
   }
 
-ET_FORALL_REAL_TYPES(GENERATE_TEST)
+ET_FORALL_REALH_TYPES(GENERATE_TEST)
+
+TEST_F(OpFullOutTest, ValueOverflow) {
+  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
+    GTEST_SKIP() << "ATen kernel doesn't handle overflow";
+  }
+  TensorFactory<ScalarType::Byte> tf;
+
+  std::vector<int64_t> sizes_int64_t_vec = {2, 3};
+  std::vector<int32_t> sizes_in32_t_vec = {2, 3};
+  auto sizes = IntArrayRef(sizes_int64_t_vec.data(), sizes_int64_t_vec.size());
+
+  Tensor out = tf.zeros(sizes_in32_t_vec);
+
+  op_full_out(sizes, 1000, out);
+}
+
+TEST_F(OpFullOutTest, HalfSupport) {
+  TensorFactory<ScalarType::Half> tf;
+
+  std::vector<int64_t> sizes_int64_t_vec = {2, 3};
+  std::vector<int32_t> sizes_in32_t_vec = {2, 3};
+  auto sizes = IntArrayRef(sizes_int64_t_vec.data(), sizes_int64_t_vec.size());
+
+  // Boolean Scalar
+  Tensor out = tf.zeros(sizes_in32_t_vec);
+  op_full_out(sizes, true, out);
+  EXPECT_TENSOR_EQ(out, tf.ones(sizes_in32_t_vec));
+
+  // Integral Scalar
+  out = tf.zeros(sizes_in32_t_vec);
+  op_full_out(sizes, 1, out);
+  EXPECT_TENSOR_EQ(out, tf.ones(sizes_in32_t_vec));
+
+  // Floating Point Scalar
+  out = tf.zeros(sizes_in32_t_vec);
+  op_full_out(sizes, 3.1415926535, out);
+  EXPECT_TENSOR_EQ(out, tf.full(sizes_in32_t_vec, 3.1415926535));
+}
+
+TEST_F(OpFullOutTest, ZeroDim) {
+  TensorFactory<ScalarType::Half> tf;
+
+  std::vector<int64_t> sizes_int64_t_vec = {};
+  std::vector<int32_t> sizes_in32_t_vec = {};
+  auto sizes = IntArrayRef(sizes_int64_t_vec.data(), sizes_int64_t_vec.size());
+
+  // Boolean Scalar
+  Tensor out = tf.zeros(sizes_in32_t_vec);
+  op_full_out(sizes, true, out);
+  EXPECT_TENSOR_EQ(out, tf.ones(sizes_in32_t_vec));
+}