From 6070f8e9abf204ce8e445d0f59110a7cfea09c6f Mon Sep 17 00:00:00 2001
From: Manuel Candales <mcandales@fb.com>
Date: Mon, 2 Oct 2023 17:25:33 -0700
Subject: [PATCH] fix resizing scalar tensor with empty sizes (#548)

Summary:

Resizing scalar tensor providing empty sizes should run successfully. Resizing logic assumed the sizes were non-empty, resulting in a crash

Reviewed By: JacobSzwejbka, dbort

Differential Revision: D49789758
---
 .../exec_aten/util/test/tensor_util_test.cpp  |  8 +++++++
 runtime/core/portable_type/tensor_impl.cpp    | 12 ++++++++++
 .../portable_type/test/tensor_impl_test.cpp   | 24 +++++++++++++++++++
 3 files changed, 44 insertions(+)

diff --git a/runtime/core/exec_aten/util/test/tensor_util_test.cpp b/runtime/core/exec_aten/util/test/tensor_util_test.cpp
index 62c723c8afc..e00032f6a2f 100644
--- a/runtime/core/exec_aten/util/test/tensor_util_test.cpp
+++ b/runtime/core/exec_aten/util/test/tensor_util_test.cpp
@@ -543,3 +543,11 @@ TEST_F(TensorUtilTest, TensorIsContiguous) {
   EXPECT_TRUE(tensor_is_contiguous(c));
   EXPECT_TRUE(tensor_is_contiguous(d));
 }
+
+TEST_F(TensorUtilTest, ResizeZeroDimTensor) {
+  using namespace torch::executor;
+  Tensor a = tf_float_.ones({});
+
+  EXPECT_EQ(resize_tensor(a, {}), Error::Ok);
+  EXPECT_EQ(a.dim(), 0);
+}
diff --git a/runtime/core/portable_type/tensor_impl.cpp b/runtime/core/portable_type/tensor_impl.cpp
index 068249a8a47..21c2061129b 100644
--- a/runtime/core/portable_type/tensor_impl.cpp
+++ b/runtime/core/portable_type/tensor_impl.cpp
@@ -113,6 +113,18 @@ Error TensorImpl::internal_resize_contiguous(ArrayRef<SizesType> new_sizes) {
       dim_,
       new_sizes.size());
 
+  // Kernels don't check that the provided out tensors have the right size.
+  // Instead they always attempt to resize the out tensor to the right size,
+  // even when the out tensor already had the right size. Therefore, if we call
+  // an op with inputs that will produce a zero-dimensional output, and the out
+  // tensor that we pass has non-STATIC dynamism, then we will end up here.
+  // Since we have already checked above that the out tensor has the right
+  // number of dimensions, it must be that the provided out tensor has zero
+  // rank, therefore it already has the right size and we should just return.
+  if (dim_ == 0) {
+    return Error::Ok;
+  }
+
   // Can only resize a StaticShape Tensor to the same size
   if (shape_dynamism_ == TensorShapeDynamism::STATIC) {
     for (int i = 0; i < new_sizes.size(); i++) {
diff --git a/runtime/core/portable_type/test/tensor_impl_test.cpp b/runtime/core/portable_type/test/tensor_impl_test.cpp
index f93815edd46..b1177766f3e 100644
--- a/runtime/core/portable_type/test/tensor_impl_test.cpp
+++ b/runtime/core/portable_type/test/tensor_impl_test.cpp
@@ -181,6 +181,30 @@ TEST_F(TensorImplTest, TestSetSizesContigUpperBounded) {
   ET_EXPECT_DEATH(t.set_sizes_contiguous({new_sizes_4, 1}), "");
 }
 
+TEST_F(TensorImplTest, TestZeroDimSetEmptySizesContig) {
+  SizesType sizes[0] = {};
+  DimOrderType dim_order[0] = {};
+  StridesType strides[0] = {};
+  float data[1] = {1.0};
+  TensorImpl t(
+      ScalarType::Float,
+      0,
+      sizes,
+      data,
+      dim_order,
+      strides,
+      TensorShapeDynamism::DYNAMIC_BOUND);
+
+  ArrayRef<SizesType> new_sizes_empty{};
+  // Can resize with empty sizes
+  t.set_sizes_contiguous(new_sizes_empty);
+  EXPECT_EQ(t.dim(), 0);
+
+  SizesType new_sizes_1[1] = {1};
+  // Can't change rank of tensor
+  ET_EXPECT_DEATH(t.set_sizes_contiguous({new_sizes_1, 1}), "");
+}
+
 TEST_F(TensorImplTest, TestWriteRead) {
   SizesType sizes[1] = {1};
   DimOrderType dim_order[1] = {0};