fix(//core/conversion/converters/impl/reduce): Adds support for multiple

reduction dimensions

Signed-off-by: Naren Dasan <naren@narendasan.com>
Signed-off-by: Naren Dasan <narens@nvidia.com>

Author: narendasan

core/conversion/converters/impl/reduce.cpp

@@ -1,3 +1,4 @@
+#include <bitset>
 #include "core/util/prelude.h"
 #include "core/conversion/converters/converters.h"
 
@@ -22,25 +23,36 @@ auto reduced_registrations = RegisterNodeConversionPatterns()
             TRTORCH_CHECK(mean_layer, "Unable to create mean layer from node: " << *n);
 
             mean_layer->setName(util::node_info(n).c_str());
-            ctx->AssociateValueAndTensor(n->outputs()[0], mean_layer->getOutput(0));
+            auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], mean_layer->getOutput(0));
+
+            LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
             return true;
         }
     }).pattern({
-        "aten::mean.dim(Tensor self, int[1] dim, bool keepdim=False, *, int? dtype=None) -> (Tensor)",
+        "aten::mean.dim(Tensor self, int[] dim, bool keepdim=False, *, int? dtype=None) -> (Tensor)",
         [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
             auto in_tensor = args[0].ITensor();
-            auto dim = args[1].unwrapToIntList();
-            auto keepdim = args[2].unwrapToBool();
+            auto dims = args[1].unwrapToIntList();
+            LOG_DEBUG("Dim to reduce:" << util::toDims(dims)); // Some abuse of toDim but just for debug info
 
-            uint32_t axis_mask = 1 << dim[0];
+            uint32_t axis_mask = 0;
+            for (int d = 0; d < dims.size(); d++) {
+                axis_mask |= 1 << dims[d];
+            }
+            LOG_DEBUG("Axis Mask" << std::bitset<32>(axis_mask));
+
+            auto keepdim = args[2].unwrapToBool();
+            LOG_DEBUG("Keep dims :" << keepdim);
 
             LOG_WARNING("Mean converter disregards dtype");
             auto mean_layer = ctx->net->addReduce(*in_tensor, nvinfer1::ReduceOperation::kAVG, axis_mask, keepdim);
 
             TRTORCH_CHECK(mean_layer, "Unable to create mean layer from node: " << *n);
 
             mean_layer->setName(util::node_info(n).c_str());
-            ctx->AssociateValueAndTensor(n->outputs()[0], mean_layer->getOutput(0));
+            auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], mean_layer->getOutput(0));
+
+            LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
             return true;
         }
     });

tests/core/converters/test_reduce.cpp

@@ -59,7 +59,32 @@ TEST(Converters, ATenMeanRowConvertsCorrectly) {
     auto g = std::make_shared<torch::jit::Graph>();
     torch::jit::script::parseIR(graph, &*g);
 
-    auto in = at::randint(-5, 5, {4, 4}, at::kCUDA);
+    auto in = at::randint(-5, 5, {4, 4, 4}, 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::RunGraphEngine(g, params, {in});
+
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0]));
+}
+
+TEST(Converters, ATenMeanMultiDimsConvertsCorrectly) {
+    const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %1 : int = prim::Constant[value=0]()
+        %2 : int = prim::Constant[value=1]()
+        %3 : int[] = prim::ListConstruct(%1, %2)
+        %4 : bool = prim::Constant[value=0]()
+        %5 : None = prim::Constant()
+        %6 : Tensor = aten::mean(%0, %3, %4, %5)
+        return (%6))IR";
+
+    auto g = std::make_shared<torch::jit::Graph>();
+    torch::jit::script::parseIR(graph, &*g);
+
+    auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA);
     auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
     auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});