Fix the coredump described by #106702 (#108002)

FFFrog · pytorchmergebot · commit 78810d78e82f · 2023-08-29T17:14:16.000Z
Fixes #106702 and add some tests As shown by [maxUnpool1d](https://pytorch.org/docs/master/generated/torch.nn.MaxUnpool1d)(`MaxUnpool2d`, `MaxUnpool3d` also), `Input` and `Output` support `(N,C,*)` or `(C,*)`, but the c++ api currently supports the former, and the latter will cause a coredump. Pull Request resolved: #108002 Approved by: https://github.com/albanD
diff --git a/test/cpp/api/functional.cpp b/test/cpp/api/functional.cpp
@@ -891,6 +891,16 @@ TEST_F(FunctionalTest, MaxUnpool1d) {
       y, torch::tensor({{{0, 2, 0, 4, 5, 0, 0, 0, 0}}}, torch::kFloat)));
   ASSERT_EQ(y.sizes(), std::vector<int64_t>({1, 1, 9}));
 
+  x = torch::tensor(
+      {{2, 4, 5}}, torch::dtype(torch::kFloat).requires_grad(true));
+  indices = torch::tensor({{1, 3, 4}}, torch::kLong);
+  y = F::max_unpool1d(x, indices, F::MaxUnpool1dFuncOptions(3));
+
+  ASSERT_EQ(y.ndimension(), 2);
+  ASSERT_TRUE(torch::allclose(
+      y, torch::tensor({{0, 2, 0, 4, 5, 0, 0, 0, 0}}, torch::kFloat)));
+  ASSERT_EQ(y.sizes(), std::vector<int64_t>({1, 9}));
+
   x = torch::tensor(
       {{{2, 4, 5}}}, torch::dtype(torch::kFloat).requires_grad(true));
   indices = torch::tensor({{{1, 3, 4}}}, torch::kLong);
@@ -945,6 +955,34 @@ TEST_F(FunctionalTest, MaxUnpool2d) {
              {0, 46, 0, 48, 49}}}},
           torch::kFloat)));
   ASSERT_EQ(y.sizes(), std::vector<int64_t>({2, 1, 5, 5}));
+
+  indices = torch::tensor(
+      {{{6, 8, 9}, {16, 18, 19}, {21, 23, 24}},
+       {{6, 8, 9}, {16, 18, 19}, {21, 23, 24}}},
+      torch::kLong);
+  x = torch::tensor(
+      {{{6, 8, 9}, {16, 18, 19}, {21, 23, 24}},
+       {{31, 33, 34}, {41, 43, 44}, {46, 48, 49}}},
+      torch::dtype(torch::kFloat).requires_grad(true));
+  y = F::max_unpool2d(
+      x, indices, F::MaxUnpool2dFuncOptions(3).stride(2).padding(1));
+
+  ASSERT_EQ(y.dim(), 3);
+  ASSERT_TRUE(torch::allclose(
+      y,
+      torch::tensor(
+          {{{0, 0, 0, 0, 0},
+            {0, 6, 0, 8, 9},
+            {0, 0, 0, 0, 0},
+            {0, 16, 0, 18, 19},
+            {0, 21, 0, 23, 24}},
+           {{0, 0, 0, 0, 0},
+            {0, 31, 0, 33, 34},
+            {0, 0, 0, 0, 0},
+            {0, 41, 0, 43, 44},
+            {0, 46, 0, 48, 49}}},
+          torch::kFloat)));
+  ASSERT_EQ(y.sizes(), std::vector<int64_t>({2, 5, 5}));
 }
 
 TEST_F(FunctionalTest, MaxUnpool3d) {
@@ -962,6 +1000,21 @@ TEST_F(FunctionalTest, MaxUnpool3d) {
              {{0, 0, 0}, {0, 0, 0}, {0, 0, 26}}}}},
           torch::kFloat)));
   ASSERT_EQ(y.sizes(), std::vector<int64_t>({1, 1, 3, 3, 3}));
+
+  indices = torch::tensor({{{{26}}}}, torch::kLong);
+  x = torch::tensor(
+      {{{{26}}}}, torch::dtype(torch::kFloat).requires_grad(true));
+  y = F::max_unpool3d(x, indices, F::MaxUnpool3dFuncOptions(3));
+
+  ASSERT_EQ(y.dim(), 4);
+  ASSERT_TRUE(torch::allclose(
+      y,
+      torch::tensor(
+          {{{{0, 0, 0}, {0, 0, 0}, {0, 0, 0}},
+            {{0, 0, 0}, {0, 0, 0}, {0, 0, 0}},
+            {{0, 0, 0}, {0, 0, 0}, {0, 0, 26}}}},
+          torch::kFloat)));
+  ASSERT_EQ(y.sizes(), std::vector<int64_t>({1, 3, 3, 3}));
 }
 
 TEST_F(FunctionalTest, ELU) {
diff --git a/torch/csrc/api/include/torch/nn/functional/pooling.h b/torch/csrc/api/include/torch/nn/functional/pooling.h
@@ -637,7 +637,9 @@ inline std::vector<int64_t> _unpool_output_size(
   std::vector<int64_t> default_size;
   for (const auto d : c10::irange(kernel_size.size())) {
     default_size.push_back(
-        (input_size[d + 2] - 1) * stride[d] + kernel_size[d] - 2 * padding[d]);
+        (input_size[input_size.size() - kernel_size.size() + d] - 1) *
+            stride[d] +
+        kernel_size[d] - 2 * padding[d]);
   }
   if (!output_size) {
     return default_size;
@@ -691,8 +693,8 @@ inline Tensor max_unpool1d(
       _unpool_output_size(input, kernel_size, stride, padding, output_size);
   output_size_.push_back(1);
   return torch::max_unpool2d(
-             input.unsqueeze(3), indices.unsqueeze(3), output_size_)
-      .squeeze(3);
+             input.unsqueeze(-1), indices.unsqueeze(-1), output_size_)
+      .squeeze(-1);
 }
 } // namespace detail
 #endif /* DOXYGEN_SHOULD_SKIP_THIS */
