support higher order grad for log

apache · May 18, 2019 · 9d1ddb3 · 9d1ddb3
1 parent 3d6646a
commit 9d1ddb3
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/src/operator/tensor/elemwise_unary_op_basic.cc b/src/operator/tensor/elemwise_unary_op_basic.cc
@@ -989,7 +989,16 @@ The storage type of ``log`` output is always dense
 
 )code" ADD_FILELINE)
 .set_attr<FCompute>("FCompute<cpu>", UnaryOp::LogCompute<cpu, mshadow_op::log>)
-.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_log"});
+.set_attr<nnvm::FGradient>("FGradient",
+  [](const nnvm::NodePtr& n, const std::vector<nnvm::NodeEntry>& ograds) {
+    auto x_grad = MakeNode("reciprocal", n->attrs.name + "_mid_x_grad",
+                            {n->inputs[0]}, nullptr, &n);
+    auto in_grad = MakeNode("elemwise_mul", n->attrs.name + "_backward",
+                            {ograds[0], nnvm::NodeEntry{x_grad, 0, 0}}, nullptr, &n);
+    std::vector<nnvm::NodeEntry> ret;
+    ret.emplace_back(nnvm::NodeEntry{in_grad, 0, 0});
+    return ret;
+  });
 
 // log10
 MXNET_OPERATOR_REGISTER_UNARY_WITH_SPARSE_DR(log10, cpu, mshadow_op::log10)