Fix & cleanup op scalar_tensor (#702)

Summary:
Pull Request resolved: #702

Resize out tensor
ghstack-source-id: 203341576
exported-using-ghexport

Reviewed By: SS-JIA

Differential Revision: D49792556

fbshipit-source-id: f2b276b0b986f9d1e0d5fc5720149454a5a3223a

Author: manuelcandales

kernels/portable/cpu/op_scalar_tensor.cpp

@@ -8,10 +8,6 @@
 
 #include <executorch/kernels/portable/cpu/scalar_utils.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
-#include <executorch/runtime/platform/assert.h>
-
-#include <cstdint>
-#include <cstring>
 
 namespace torch {
 namespace executor {
@@ -20,16 +16,17 @@ namespace native {
 Tensor& scalar_tensor_out(RuntimeContext& ctx, const Scalar& s, Tensor& out) {
   (void)ctx;
 
-  ET_CHECK_MSG(out.numel() == 1, "Output tensor must have only one element");
+  ET_KERNEL_CHECK(
+      ctx, resize_tensor(out, {}) == Error::Ok, InvalidArgument, out);
 
   ScalarType s_type = utils::get_scalar_dtype(s);
   ScalarType out_type = out.scalar_type();
 
-  ET_SWITCH_REAL_TYPES_AND(Bool, out_type, ctx, "scalar_tensor", CTYPE, [&]() {
-    ET_SWITCH_SCALAR_OBJ_TYPES(s_type, ctx, "scalar_tensor", CTYPE_S, [&]() {
+  ET_SWITCH_REAL_TYPES_AND(Bool, out_type, ctx, __func__, CTYPE, [&]() {
+    ET_SWITCH_SCALAR_OBJ_TYPES(s_type, ctx, __func__, CTYPE_S, [&]() {
       CTYPE_S val_s;
       ET_EXTRACT_SCALAR(s, val_s);
-      out.mutable_data_ptr<CTYPE>()[0] = val_s;
+      out.mutable_data_ptr<CTYPE>()[0] = convert<CTYPE, CTYPE_S>(val_s);
     });
   });
 

kernels/test/op_scalar_tensor_test.cpp

@@ -55,38 +55,29 @@ void test_scalar_tensor_out_1d(CTYPE value) {
   TensorFactory<DTYPE> tf;
 
   std::vector<int32_t> sizes{1};
-  Tensor expected = tf.make(sizes, /*data=*/{value});
-
   Tensor out = tf.ones(sizes);
-  op_scalar_tensor_out(value, out);
 
-  EXPECT_TENSOR_EQ(out, expected);
+  ET_EXPECT_KERNEL_FAILURE(op_scalar_tensor_out(value, out));
 }
 
 template <typename CTYPE, ScalarType DTYPE>
 void test_scalar_tensor_out_2d(CTYPE value) {
   TensorFactory<DTYPE> tf;
 
   std::vector<int32_t> sizes{1, 1};
-  Tensor expected = tf.make(sizes, /*data=*/{value});
-
   Tensor out = tf.ones(sizes);
-  op_scalar_tensor_out(value, out);
 
-  EXPECT_TENSOR_EQ(out, expected);
+  ET_EXPECT_KERNEL_FAILURE(op_scalar_tensor_out(value, out));
 }
 
 template <typename CTYPE, ScalarType DTYPE>
 void test_scalar_tensor_out_3d(CTYPE value) {
   TensorFactory<DTYPE> tf;
 
   std::vector<int32_t> sizes{1, 1, 1};
-  Tensor expected = tf.make(sizes, /*data=*/{value});
-
   Tensor out = tf.ones(sizes);
-  op_scalar_tensor_out(value, out);
 
-  EXPECT_TENSOR_EQ(out, expected);
+  ET_EXPECT_KERNEL_FAILURE(op_scalar_tensor_out(value, out));
 }
 
 #define GENERATE_TEST(ctype, dtype)                                    \