fix(aten::flatten): Fixes dynamic shape for flatten

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

Diff for: core/conversion/converters/impl/shuffle.cpp

@@ -17,7 +17,12 @@ static auto shuffle_registrations TRTORCH_UNUSED = RegisterNodeConversionPattern
       auto start_dim = args[1].unwrapToInt();
       auto end_dim = args[2].unwrapToInt();
       auto in_shape = util::toVec(in->getDimensions());
-      auto out_shape = torch::flatten(torch::rand(in_shape), start_dim, end_dim).sizes();
+      std::vector<int64_t> out_shape;
+      if (ctx->input_is_dynamic) {
+        out_shape = std::vector<int64_t>({in_shape[0], -1});
+      } else {
+        out_shape = torch::flatten(torch::rand(in_shape), start_dim, end_dim).sizes().vec();
+      }
 
       auto shuffle = ctx->net->addShuffle(*in);
       TRTORCH_CHECK(shuffle, "Unable to create shuffle layer from node: " << *n);
@@ -33,7 +38,12 @@ static auto shuffle_registrations TRTORCH_UNUSED = RegisterNodeConversionPattern
     [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
       auto in = args[0].ITensor();
       auto in_shape = util::toVec(in->getDimensions());
-      auto new_shape = torch::reshape(torch::rand(in_shape), args[1].unwrapToIntList().vec()).sizes();
+      std::vector<int64_t> new_shape;
+      if (ctx->input_is_dynamic) {
+        TRTORCH_THROW_ERROR("Resize is currently not support in dynamic input shape compilation");
+      } else {
+        new_shape = torch::reshape(torch::rand(in_shape), args[1].unwrapToIntList().vec()).sizes().vec();
+      }
 
       auto shuffle = ctx->net->addShuffle(*in);
       TRTORCH_CHECK(shuffle, "Unable to create shuffle layer from node: " << *n);

Diff for: tests/core/converters/test_shuffle.cpp

@@ -6,26 +6,26 @@
 
 // TODO: IR Parser doesnt work well with neg numbers
 TEST(Converters, ATenFlattenConvertsCorrectly) {
-    const auto graph = R"IR(
-      graph(%0 : Tensor):
-        %1 : int = prim::Constant[value=0]()
-        %2 : int = prim::Constant[value=1]()
-        %3 : Tensor = aten::flatten(%0, %1, %2)
-        return (%3))IR";
-
-    auto g = std::make_shared<torch::jit::Graph>();
-    torch::jit::parseIR(graph, &*g);
+  const auto graph = R"IR(
+    graph(%0 : Tensor):
+      %1 : int = prim::Constant[value=0]()
+      %2 : int = prim::Constant[value=1]()
+      %3 : Tensor = aten::flatten(%0, %1, %2)
+      return (%3))IR";
 
-    auto in = at::randint(0, 5, {2, 3}, {at::kCUDA});
-    auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
-    auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
 
-    in = at::clone(in);
-    params = trtorch::core::conversion::get_named_params(g->inputs(), {});
-    auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
-    auto trt = trt_results[0].reshape_as(jit_results[0]);
+  auto in = at::randint(0, 5, {2, 3}, {at::kCUDA});
+  auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
 
-    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+  in = at::clone(in);
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
+  auto trt = trt_results[0].reshape_as(jit_results[0]);
+
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
 }
 
 // TODO: IR Parser doesnt work well with neg numbers
@@ -164,4 +164,27 @@ TEST(Converters, ATenPermute5DConvertsCorrectly) {
     auto trt = trt_results[0].reshape_as(jit_results[0]);
 
     ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
-}
+}
+
+TEST(Converters, ATenFlattenConvertsCorrectlyWithDynamicInput) {
+  const auto graph = R"IR(
+    graph(%0 : Tensor):
+      %1 : int = prim::Constant[value=0]()
+      %2 : int = prim::Constant[value=1]()
+      %3 : Tensor = aten::flatten(%0, %1, %2)
+      return (%3))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::randint(0, 5, {2, 3}, {at::kCUDA});
+  auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
+
+  in = at::clone(in);
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto trt_results = trtorch::tests::util::RunGraphEngineDynamic(g, params, {in});
+  auto trt = trt_results[0].reshape_as(jit_results[0]);
+
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}