support index=-1, and fix bug when nDims=2 or 3 in softmax

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

Author: inocsin

core/conversion/converters/impl/softmax.cpp

@@ -26,12 +26,17 @@ static auto softmax_registrations TRTORCH_UNUSED = RegisterNodeConversionPattern
        }
 
        int64_t dim = args[1].IValue()->toInt();
+       LOG_DEBUG("Softmax original dim " << dim);
+       if (dim == -1) {
+         dim = shape.size() - 1;
+       }
+       LOG_DEBUG("Softmax converted dim " << dim);
        auto softmax = ctx->net->addSoftMax(*in);
 
        TRTORCH_CHECK(softmax, "Unable to create softmax layer from node: " << *n);
        LOG_DEBUG("Disregarding dtype argument");
 
-       if (shape.size() > 3) {
+       if (shape.size() > 1) {
          softmax->setAxes(1 << (dim));
        } else {
          // When there is no batch dimension

tests/core/conversion/converters/test_softmax.cpp

@@ -76,3 +76,29 @@ TEST(Converters, ATenSoftmaxNDConvertsCorrectlyAbove3DIndex) {
 
   ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
 }
+
+TEST(Converters, ATenSoftmaxNDConvertsCorrectlyNegtiveIndex) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %1 : None = prim::Constant()
+        %2 : int = prim::Constant[value=-1]()
+        %3 : Tensor = aten::softmax(%0, %2, %1)
+        return (%3))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::randint(0, 5, {1, 2, 2, 2, 2}, {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_as(jit_results[0]);
+
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}