Make TensorPtr constructor check the data dize matches the shape. (pytorch#13591)

Summary: .

Differential Revision: D80764139

Author: shoumikhin

extension/tensor/tensor_ptr.cpp

@@ -148,10 +148,10 @@ TensorPtr make_tensor_ptr(
     executorch::aten::ScalarType type,
     executorch::aten::TensorShapeDynamism dynamism) {
   ET_CHECK_MSG(
-      data.size() >=
+      data.size() ==
           executorch::aten::compute_numel(sizes.data(), sizes.size()) *
               executorch::aten::elementSize(type),
-      "Data size is smaller than required by sizes and scalar type.");
+      "Data size does not match tensor size.");
   auto data_ptr = data.data();
   return make_tensor_ptr(
       std::move(sizes),

extension/tensor/tensor_ptr.h

@@ -106,6 +106,10 @@ inline TensorPtr make_tensor_ptr(
     executorch::aten::ScalarType type = deduced_type,
     executorch::aten::TensorShapeDynamism dynamism =
         executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND) {
+  ET_CHECK_MSG(
+      data.size() ==
+          executorch::aten::compute_numel(sizes.data(), sizes.size()),
+      "Data size does not match tensor size.");
   if (type != deduced_type) {
     ET_CHECK_MSG(
         runtime::canCast(deduced_type, type),

extension/tensor/test/tensor_ptr_test.cpp

@@ -784,16 +784,30 @@ TEST_F(TensorPtrTest, TensorUint8BufferTooSmallExpectDeath) {
       { auto tensor = make_tensor_ptr({2, 2}, std::move(data)); }, "");
 }
 
-TEST_F(TensorPtrTest, TensorUint8BufferTooLarge) {
+TEST_F(TensorPtrTest, TensorUint8BufferTooLargeExpectDeath) {
   std::vector<uint8_t> data(
-      4 * executorch::aten::elementSize(executorch::aten::ScalarType::Float));
-  auto tensor = make_tensor_ptr({2, 2}, std::move(data));
+      5 * executorch::aten::elementSize(executorch::aten::ScalarType::Float));
+  ET_EXPECT_DEATH({ auto _ = make_tensor_ptr({2, 2}, std::move(data)); }, "");
+}
 
-  EXPECT_EQ(tensor->dim(), 2);
-  EXPECT_EQ(tensor->size(0), 2);
-  EXPECT_EQ(tensor->size(1), 2);
-  EXPECT_EQ(tensor->strides()[0], 2);
-  EXPECT_EQ(tensor->strides()[1], 1);
+TEST_F(TensorPtrTest, VectorFloatTooSmallExpectDeath) {
+  std::vector<float> data(9, 1.f);
+  ET_EXPECT_DEATH({ auto _ = make_tensor_ptr({2, 5}, std::move(data)); }, "");
+}
+
+TEST_F(TensorPtrTest, VectorFloatTooLargeExpectDeath) {
+  std::vector<float> data(11, 1.f);
+  ET_EXPECT_DEATH({ auto _ = make_tensor_ptr({2, 5}, std::move(data)); }, "");
+}
+
+TEST_F(TensorPtrTest, VectorIntToFloatCastTooSmallExpectDeath) {
+  std::vector<int32_t> data(9, 1);
+  ET_EXPECT_DEATH({ auto _ = make_tensor_ptr({2, 5}, std::move(data)); }, "");
+}
+
+TEST_F(TensorPtrTest, VectorIntToFloatCastTooLargeExpectDeath) {
+  std::vector<int32_t> data(11, 1);
+  ET_EXPECT_DEATH({ auto _ = make_tensor_ptr({2, 5}, std::move(data)); }, "");
 }
 
 TEST_F(TensorPtrTest, StridesAndDimOrderMustMatchSizes) {