From 8eca88c05bb46c24a0c101426e4c8c8f7f488548 Mon Sep 17 00:00:00 2001
From: Tianqi Chen <tqchen@users.noreply.github.com>
Date: Wed, 19 Sep 2018 12:49:20 -0700
Subject: [PATCH] [NNVM] Recover reduction behavir, fix CI (#1740)

---
 nnvm/src/top/tensor/reduce.cc |  5 +++++
 topi/include/topi/reduction.h | 11 +++++++++++
 topi/src/topi.cc              |  2 ++
 3 files changed, 18 insertions(+)

diff --git a/nnvm/src/top/tensor/reduce.cc b/nnvm/src/top/tensor/reduce.cc
index 527a6a5abd743..91d2ea7202b84 100644
--- a/nnvm/src/top/tensor/reduce.cc
+++ b/nnvm/src/top/tensor/reduce.cc
@@ -3,6 +3,9 @@
  * \file reduce.cc
  * \brief reduce operator.
  */
+// Enforce TOPI to use old behavior that reduces to at least 1d
+#define TOPI_REDUCE_ATLEAST1D 1
+
 #include <nnvm/op.h>
 #include <nnvm/node.h>
 #include <nnvm/op_attr_types.h>
@@ -17,6 +20,8 @@
 #include "topi/reduction.h"
 #include "topi/transform.h"
 
+static_assert(TOPI_REDUCE_ATLEAST1D, "need to use legacy reduce behavior");
+
 namespace nnvm {
 namespace top {
 using namespace tvm;
diff --git a/topi/include/topi/reduction.h b/topi/include/topi/reduction.h
index ccc85e96643ee..d68b9b3904195 100644
--- a/topi/include/topi/reduction.h
+++ b/topi/include/topi/reduction.h
@@ -20,6 +20,14 @@
 #include "topi/detail/constant_utils.h"
 #include "tvm/tvm.h"
 
+/*!
+ * \brief macro flag to enable some legacy behavior which requires
+ * reduction result to be at least 1d.
+ */
+#ifndef TOPI_REDUCE_ATLEAST1D
+#define TOPI_REDUCE_ATLEAST1D 0
+#endif
+
 namespace topi {
 using namespace tvm;
 
@@ -96,6 +104,9 @@ inline Array<Expr> MakeReduceTargetShape(const std::vector<int>& real_axis,
       }
     }
   }
+  if (target_shape.size() == 0 && TOPI_REDUCE_ATLEAST1D) {
+    target_shape.push_back(1);
+  }
   return target_shape;
 }
 
diff --git a/topi/src/topi.cc b/topi/src/topi.cc
index cac3545a75a2e..fef2487e67702 100644
--- a/topi/src/topi.cc
+++ b/topi/src/topi.cc
@@ -3,6 +3,8 @@
 * \brief Registration of TVM operators and schedules
 * \file topi.cc
 */
+#define TOPI_REDUCE_ATLEAST1D 0
+
 #include <tvm/runtime/packed_func.h>
 #include <tvm/runtime/module.h>
 #include <tvm/runtime/registry.h>