Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
[NNVM] Recover reduction behavir, fix CI (#1740)
  • Loading branch information
tqchen committed Sep 19, 2018
1 parent ac7b420 commit 543c424
Show file tree
Hide file tree
Showing 3 changed files with 18 additions and 0 deletions.
5 changes: 5 additions & 0 deletions nnvm/src/top/tensor/reduce.cc
Expand Up @@ -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>
Expand All @@ -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;
Expand Down
11 changes: 11 additions & 0 deletions topi/include/topi/reduction.h
Expand Up @@ -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;

Expand Down Expand Up @@ -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;
}

Expand Down
2 changes: 2 additions & 0 deletions topi/src/topi.cc
Expand Up @@ -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>
Expand Down

0 comments on commit 543c424

Please sign in to comment.