feat(aten::sum): Allow for negative indices less than -1

Signed-off-by: Naren Dasan <naren@narendasan.com> Signed-off-by: Naren Dasan <narens@nvidia.com>
pytorch · Feb 8, 2021 · 769bbc9 · 769bbc9
1 parent 3e7cf8e
commit 769bbc9
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 2 deletions.
diff --git a/core/conversion/converters/impl/reduce.cpp b/core/conversion/converters/impl/reduce.cpp
@@ -85,8 +85,8 @@ auto reduce_registrations TRTORCH_UNUSED =
                LOG_DEBUG("InDims " << in_dims); // Some abuse of toDim but just for debug info
                LOG_DEBUG(
                    "Dim to reduce(original):" << util::toDims(dims)); // Some abuse of toDim but just for debug info
-               for (int i = 0; i < dims.size(); i++) {
-                 auto dim_val = dims[i] == -1 ? (in_dims.size() - 1) : dims[i];
+               for (size_t i = 0; i < dims.size(); i++) {
+                 auto dim_val = dims[i] < 0 ? (in_dims.size() + dims[i]) : dims[i];
                  calculated_dims.push_back(dim_val);
                }
 

diff --git a/tests/core/conversion/converters/test_reduce.cpp b/tests/core/conversion/converters/test_reduce.cpp
@@ -134,6 +134,58 @@ converts_keepdims_correctly(mean, Mean);
 
 #undef converts_keepdims_correctly
 
+TEST(Converters, ATenSumDimNegOneIndexConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%0 : Tensor):
+      %1 : int = prim::Constant[value=-1]()
+      %2 : int[] = prim::ListConstruct(%1)
+      %3 : bool = prim::Constant[value=0]()
+      %4 : None = prim::Constant()
+      %5 : Tensor = aten::sum(%0, %2, %3, %4)
+      return (%5))IR";
+  auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA);
+  test_body(graph, in);
+}
+
+TEST(Converters, ATenSumDimNegOneIndexKeepDimsConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%0 : Tensor):
+      %1 : int = prim::Constant[value=-1]()
+      %2 : int[] = prim::ListConstruct(%1)
+      %3 : bool = prim::Constant[value=1]()
+      %4 : None = prim::Constant()
+      %5 : Tensor = aten::sum(%0, %2, %3, %4)
+      return (%5))IR";
+  auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA);
+  test_body(graph, in);
+}
+
+TEST(Converters, ATenSumDimNegIndexConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%0 : Tensor):
+      %1 : int = prim::Constant[value=-2]()
+      %2 : int[] = prim::ListConstruct(%1)
+      %3 : bool = prim::Constant[value=0]()
+      %4 : None = prim::Constant()
+      %5 : Tensor = aten::sum(%0, %2, %3, %4)
+      return (%5))IR";
+  auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA);
+  test_body(graph, in);
+}
+
+TEST(Converters, ATenSumDimNegIndexKeepDimsConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%0 : Tensor):
+      %1 : int = prim::Constant[value=-2]()
+      %2 : int[] = prim::ListConstruct(%1)
+      %3 : bool = prim::Constant[value=1]()
+      %4 : None = prim::Constant()
+      %5 : Tensor = aten::sum(%0, %2, %3, %4)
+      return (%5))IR";
+  auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA);
+  test_body(graph, in);
+}
+
 TEST(Converters, ATenProdDimConvertsCorrectly) {
   const auto graph = R"IR(
     graph(%0 : Tensor):