fix(aten::select and aten::var): Fix converters to handle negative axes

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

core/conversion/converters/impl/select.cpp

@@ -73,17 +73,23 @@ auto select_registrations TRTORCH_UNUSED =
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                     auto in = args[0].ITensorOrFreeze(ctx);
                     auto maxDim = static_cast<int64_t>(in->getDimensions().nbDims);
-                    auto axis = args[1].unwrapToInt();
-                    axis = axis < 0 ? axis + maxDim : axis;
+                    auto dim = args[1].unwrapToInt();
+                    // Handle negative axis by refering to nbDims of input Tensor
+                    dim = dim < 0 ? dim + maxDim : dim;
                     auto ind = (int32_t)args[2].unwrapToInt();
+                    // Along the specified dimension, handle negative index by subtracting along length of dimension.
+                    ind = ind < 0 ? ind + in->getDimensions().d[dim] : ind;
+                    LOG_DEBUG("Gather input dimensions: " << in->getDimensions());
+                    LOG_DEBUG("Dimension to select: " << dim);
+                    LOG_DEBUG("Index: " << ind);
 
                     // index to access needs to be an at::Tensor
                     at::Tensor indices = torch::tensor({ind}).to(torch::kI32);
                     auto const_out = tensor_to_const(ctx, indices);
 
                     // IGatherLayer takes in input tensor, the indices, and the axis
                     // of input tensor to take indices from
-                    auto gather_layer = ctx->net->addGather(*in, *const_out, axis);
+                    auto gather_layer = ctx->net->addGather(*in, *const_out, dim);
                     TRTORCH_CHECK(gather_layer, "Unable to create gather layer from node: " << *n);
                     auto out = gather_layer->getOutput(0);
 
@@ -93,7 +99,7 @@ auto select_registrations TRTORCH_UNUSED =
                       // IShuffleLayer removes redundant dimensions
                       auto shuffle_layer = ctx->net->addShuffle(*out);
                       TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
-                      shuffle_layer->setReshapeDimensions(util::squeezeDims(out->getDimensions(), axis));
+                      shuffle_layer->setReshapeDimensions(util::squeezeDims(out->getDimensions(), dim));
                       shuffle_layer->setName(util::node_info(n).c_str());
                       out = shuffle_layer->getOutput(0);
                     }

core/conversion/evaluators/eval_util.cpp

@@ -199,7 +199,7 @@ void recursiveStore(
         TRTORCH_THROW_ERROR("Found unsupported data type in arguments for aten::tensor");
       }
     } else {
-      TRTORCH_ASSERT("Found unsupported data type in arguments for aten::tensor");
+      TRTORCH_THROW_ERROR("Found unsupported data type in arguments for aten::tensor");
     }
   }
 }

core/lowering/passes/unpack_var.cpp

@@ -17,6 +17,7 @@ void UnpackVar(std::shared_ptr<torch::jit::Graph>& graph) {
       %none: None = prim::Constant()
       %false: bool = prim::Constant[value=0]()
       %0: int = prim::Constant[value=0]()
+      %dtype: int = prim::Constant[value=6]()
       %1: int = prim::Constant[value=1]()
       %sqrd: Tensor = aten::mul(%input, %input)
       %sqrdmean: Tensor = aten::mean(%sqrd, %dims, %keepdim, %none)
@@ -26,7 +27,7 @@ void UnpackVar(std::shared_ptr<torch::jit::Graph>& graph) {
       %varout : Tensor = prim::If(%unbiased)
         block0():
           %shape: int[] = aten::size(%input)
-          %shapet: Tensor = aten::tensor(%shape, %0, %none, %false)
+          %shapet: Tensor = aten::tensor(%shape, %dtype, %none, %false)
           %dim: int = prim::ListUnpack(%dims)
           %reduceddims: Tensor = aten::select(%shapet, %0, %dim)
           %numel: Tensor = aten::prod(%reduceddims, %dim, %keepdim, %none)

tests/core/conversion/converters/test_reduce.cpp

@@ -473,3 +473,29 @@ TEST(Converters, UnpackStdUnbiasedKeepDimsLowersCorrectly) {
   auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
   ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
 }
+
+TEST(Converters, UnpackVarUnbiasedNegAxisLowersCorrectly) {
+  const auto graph = R"IR(
+      graph(%x.1 : Tensor):
+        %37 : bool = prim::Constant[value=1]()
+        %53 : int[] = prim::Constant[value=[-1]]()
+        %69 : Tensor = aten::var(%x.1, %53, %37, %37)
+        return (%69))IR";
+
+  auto in = at::randint(-5, 5, {2, 20, 768}, at::kCUDA).to(at::kFloat);
+
+  auto jit_in = at::clone(in);
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto jit_results = trtorch::tests::util::RunGraph(g, params, {jit_in});
+
+  in = at::clone(in);
+  trtorch::core::lowering::passes::UnpackVar(g);
+  torch::jit::EliminateCommonSubexpression(g);
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {jit_in});
+
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
+}

tests/core/conversion/converters/test_select.cpp

@@ -85,6 +85,33 @@ TEST(Converters, ATenSelectIntDimNegativeConvertsCorrectly) {
   ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
 }
 
+TEST(Converters, ATenSelectIntNegIndexConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=0]()
+        %3 : int = prim::Constant[value=-1]()
+        %4 : Tensor = aten::select(%0, %3, %2)
+        return (%4))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+
+  torch::jit::parseIR(graph, g.get());
+
+  auto in = torch::tensor({2, 20, 768}).to(at::kFloat).to(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});
+
+  auto trt = trt_results[0].reshape(jit_results[0].sizes());
+
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}
+
 TEST(Converters, ATenSelectIntTwiceConvertsCorrectly) {
   const auto graph = R"IR(
       graph(%0 : Tensor):
@@ -437,4 +464,4 @@ TEST(Converters, ATenMaskedFillZerosConvertsCorrectly) {
   std::cout << trt_results[1].reshape_as(jit_results[0]) << std::endl;
 
   ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
-}
+}