test stack fix when dim is -1

Signed-off-by: hongwei03 <[email protected]>

Author: p517332051

tests/core/conversion/converters/test_stack.cpp

@@ -5,51 +5,72 @@
 #include "torch/csrc/jit/ir/irparser.h"
 
 TEST(Converters, ATenStackPureTensorConvertsCorrectly) {
+  auto TestATenStackPureTensorConvertsCorrectly = [](const std::string& graph) {
+    auto g = std::make_shared<torch::jit::Graph>();
+    torch::jit::parseIR(graph, g.get());
+
+    auto in1 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
+    auto in2 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
+
+    auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+    auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, in2});
+
+    params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+    auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, in2});
+
+    ASSERT_TRUE(
+        torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+  };
   const auto graph = R"IR(
       graph(%0 : Tensor,
             %1 : Tensor):
         %2 : Tensor[] = prim::ListConstruct(%0, %1)
         %3 : int = prim::Constant[value=3]()
         %4 : Tensor = aten::stack(%2, %3)
         return (%4))IR";
+  const auto graph2 = R"IR(
+      graph(%0 : Tensor,
+            %1 : Tensor):
+        %2 : Tensor[] = prim::ListConstruct(%0, %1)
+        %3 : int = prim::Constant[value=-1]()
+        %4 : Tensor = aten::stack(%2, %3)
+        return (%4))IR";
 
-  auto g = std::make_shared<torch::jit::Graph>();
-  torch::jit::parseIR(graph, g.get());
+  TestATenStackPureTensorConvertsCorrectly(graph);
+  TestATenStackPureTensorConvertsCorrectly(graph2);
+}
 
-  auto in1 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
-  auto in2 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
+TEST(Converters, ATenStackDiffTensorConvertsCorrectly) {
+  auto TestATenStackDiffTensorConvertsCorrectly = [](const std::string& graph) {
+    auto g = std::make_shared<torch::jit::Graph>();
+    torch::jit::parseIR(graph, g.get());
 
-  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
-  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, in2});
+    auto in1 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
+    auto in2 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
 
-  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
-  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1, in2});
+    auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {in2});
+    auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1});
 
-  ASSERT_TRUE(
-      torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
-}
+    params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {in2});
+    auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1});
 
-TEST(Converters, ATenStackDiffTensorConvertsCorrectly) {
+    ASSERT_TRUE(
+        torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+  };
   const auto graph = R"IR(
       graph(%0 : Tensor,
             %1 : Float(4, 4, 4, strides=[16, 4, 1])):
         %2 : Tensor[] = prim::ListConstruct(%0, %1)
         %3 : int = prim::Constant[value=1]()
         %4 : Tensor = aten::stack(%2, %3)
         return (%4))IR";
-
-  auto g = std::make_shared<torch::jit::Graph>();
-  torch::jit::parseIR(graph, g.get());
-
-  auto in1 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
-  auto in2 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
-
-  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {in2});
-  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1});
-
-  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {in2});
-  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in1});
-
-  ASSERT_TRUE(
-      torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
-}
+  const auto graph2 = R"IR(
+      graph(%0 : Tensor,
+            %1 : Float(4, 4, 4, strides=[16, 4, 1])):
+        %2 : Tensor[] = prim::ListConstruct(%0, %1)
+        %3 : int = prim::Constant[value=-1]()
+        %4 : Tensor = aten::stack(%2, %3)
+        return (%4))IR";
+  TestATenStackDiffTensorConvertsCorrectly(graph);
+  TestATenStackDiffTensorConvertsCorrectly(graph2);
+}