Merge pull request #400 from guoruoqian/select_int_fix_bug

Fix bug that when the dim is negative or bigger than one in selec…

Author: narendasan

core/conversion/converters/impl/select.cpp

@@ -68,7 +68,9 @@ auto select_registrations TRTORCH_UNUSED =
         .pattern({"aten::select.int(Tensor(a) self, int dim, int index) -> (Tensor(a))",
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                     auto in = args[0].ITensor();
+                    auto maxDim = static_cast<int64_t>(in->getDimensions().nbDims);
                     auto axis = args[1].unwrapToInt();
+                    axis = axis < 0 ? axis + maxDim : axis;
                     auto ind = (int32_t)args[2].unwrapToInt();
 
                     // index to access needs to be an at::Tensor
@@ -89,7 +91,7 @@ auto select_registrations TRTORCH_UNUSED =
                     // IShuffleLayer removes redundant dimensions
                     auto shuffle_layer = ctx->net->addShuffle(*gather_out);
                     TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
-                    shuffle_layer->setReshapeDimensions(util::unpadDims(gather_out->getDimensions()));
+                    shuffle_layer->setReshapeDimensions(util::squeezeDims(gather_out->getDimensions(), axis));
                     shuffle_layer->setName(util::node_info(n).c_str());
                     auto shuffle_out = shuffle_layer->getOutput(0);
 

tests/core/conversion/converters/test_select.cpp

@@ -31,6 +31,59 @@ TEST(Converters, ATenSelectIntConvertsCorrectly) {
   ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
 }
 
+TEST(Converters, ATenSelectIntDimIsOneConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=1]()
+        %3 : int = prim::Constant[value=0]()
+        %4 : Tensor = aten::select(%0, %2, %3)
+        return (%4))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
+
+  auto jit_in = at::clone(in);
+  auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto jit_results = trtorch::tests::util::RunGraph(g, params, {jit_in});
+
+  auto trt_in = at::clone(in);
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {trt_in});
+
+  // In order to check whether shape match that we don't do reshape.
+  // E.g. x = at::randint(1, 10, {4, 4, 4}, {at::kCUDA}), then aten::select(x, 1, 0). We should get a tensor y with
+  // shape {4, 4} instead of a tensor with shape {4, 1, 4}.
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
+}
+
+TEST(Converters, ATenSelectIntDimNegativeConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=-2]()
+        %3 : int = prim::Constant[value=0]()
+        %4 : Tensor = aten::select(%0, %2, %3)
+        return (%4))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
+
+  auto jit_in = at::clone(in);
+  auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto jit_results = trtorch::tests::util::RunGraph(g, params, {jit_in});
+
+  auto trt_in = at::clone(in);
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {trt_in});
+
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
+}
+
 TEST(Converters, ATenSelectIntTwiceConvertsCorrectly) {
   const auto graph = R"IR(
       graph(%0 : Tensor):