From 80e263f6af6f51115b2eba32d8155ca9734ebdf1 Mon Sep 17 00:00:00 2001
From: Yao Wang <kevinthesunwy@gmail.com>
Date: Thu, 8 Feb 2018 18:54:14 -0800
Subject: [PATCH] Move dense compute back to python (#364)

---
 python/nnvm/top/nn.py |  7 +++++++
 src/top/nn/nn.cc      | 15 ---------------
 2 files changed, 7 insertions(+), 15 deletions(-)

diff --git a/python/nnvm/top/nn.py b/python/nnvm/top/nn.py
index f91f77dfd..086e41185 100644
--- a/python/nnvm/top/nn.py
+++ b/python/nnvm/top/nn.py
@@ -51,6 +51,13 @@ def schedule_log_softmax(_, outs, target):
 
 
 # dense
+@reg.register_compute("dense")
+def compute_dense(attrs, inputs, _):
+    """Compute definition of dense"""
+    if attrs.get_bool("use_bias"):
+        return topi.nn.dense(inputs[0], inputs[1], bias=inputs[2])
+    return topi.nn.dense(inputs[0], inputs[1])
+
 @reg.register_schedule("dense")
 def schedule_dense(_, outs, target):
     """Schedule definition of dense"""
diff --git a/src/top/nn/nn.cc b/src/top/nn/nn.cc
index 88ed7fd58..4bb518567 100644
--- a/src/top/nn/nn.cc
+++ b/src/top/nn/nn.cc
@@ -82,21 +82,6 @@ If ``use_bias`` is set to be false, then the ``bias`` term is ignored.
 .set_attr<FListInputNames>("FListInputNames", UseBiasListInputNames<DenseParam>)
 .set_attr<FInferShape>("FInferShape", DenseInferShape)
 .set_attr<FInferType>("FInferType", ElemwiseType<-1, 1>)
-.set_attr<FTVMCompute>(
-  "FTVMCompute", [](const NodeAttrs& attrs,
-                    const Array<Tensor>& inputs,
-                    const Array<Tensor>& out_info) {
-    Tensor bias_val;
-    Tensor* bias;
-    const DenseParam& param = nnvm::get<DenseParam>(attrs.parsed);
-    if (param.use_bias) {
-      bias_val = inputs[2];
-      bias = &bias_val;
-    } else {
-      bias = nullptr;
-    }
-    return Array<Tensor>{ topi::nn::dense(inputs[0], inputs[1], bias) };
-})
 .set_attr<FGradient>(
   "FGradient", [](const NodePtr& n,
                   const std::vector<NodeEntry>& ograds) {