[incubator-mxnet] branch master updated: [MXNET-978] Higher Order Gradient Support `reciprocal`, `abs`. (#15413)

[ap...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

apeforest pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git


The following commit(s) were added to refs/heads/master by this push:
     new a3ae309  [MXNET-978] Higher Order Gradient Support `reciprocal`, `abs`. (#15413)
a3ae309 is described below

commit a3ae30979989f488cd933c2fbb6416a4e187de9d
Author: kshitij12345 <ks...@gmail.com>
AuthorDate: Mon Jul 8 02:35:09 2019 +0530

    [MXNET-978] Higher Order Gradient Support `reciprocal`, `abs`. (#15413)
    
    * add higher order support for reciprocal and abs
    
    * add relevant tests
    
    * address comments
    
    * fix extra line in tests.
    * fix missing space.
    * fix incorrect comment.
---
 src/operator/tensor/elemwise_unary_op_basic.cc  | 54 ++++++++++++++++++++++++-
 tests/python/unittest/test_higher_order_grad.py | 27 +++++++++++++
 2 files changed, 79 insertions(+), 2 deletions(-)

diff --git a/src/operator/tensor/elemwise_unary_op_basic.cc b/src/operator/tensor/elemwise_unary_op_basic.cc
index 26c7408..6da384d 100644
--- a/src/operator/tensor/elemwise_unary_op_basic.cc
+++ b/src/operator/tensor/elemwise_unary_op_basic.cc
@@ -717,7 +717,38 @@ Example::
 
 MXNET_OPERATOR_REGISTER_BINARY(_backward_reciprocal)
 .set_attr<FCompute>("FCompute<cpu>",
-  ElemwiseBinaryOp::Compute<cpu, unary_bwd<mshadow_op::reciprocal_grad> >);
+  ElemwiseBinaryOp::Compute<cpu, unary_bwd<mshadow_op::reciprocal_grad> >)
+.set_attr<nnvm::FGradient>("FGradient",
+  [](const nnvm::NodePtr& n, const std::vector<nnvm::NodeEntry>& ograds) {
+    // ograds[0]: dL/dxgrad
+    // inputs[0]: dL/dy
+    // inputs[1]: x
+    // f(x) = y = 1/x
+    // f'(x) = -1/x^2
+    // f''(x) = 2/x^3 = -2 * (f'(x) * f(x))
+
+    const std::unordered_map<std::string, std::string> args = {{"scalar", "-2.0"}};
+
+    auto dydx_mul_dldy = nnvm::NodeEntry{n};  // f'(x) * head_grads
+    auto dydx = MakeNode("elemwise_div", n->attrs.name + "_dydx",
+                         {dydx_mul_dldy, n->inputs[0]}, nullptr, &n);
+    auto fx = MakeNode("reciprocal", n->attrs.name + "_fx",
+                       {n->inputs[1]}, nullptr, &n);
+
+    auto d2ydx2_mid = MakeNode("elemwise_mul", n->attrs.name + "_d2ydx2_mid",
+                               {dydx_mul_dldy, nnvm::NodeEntry{fx}}, nullptr, &n);
+
+    auto d2ydx2 = MakeNode("_mul_scalar", n->attrs.name + "_d2ydx2",
+                           {nnvm::NodeEntry{d2ydx2_mid}}, &args, &n);
+
+    std::vector<nnvm::NodeEntry> ret;
+
+    ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad",
+                             {ograds[0], nnvm::NodeEntry{dydx}}, nullptr, &n));
+    ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad_inp",
+                             {ograds[0], nnvm::NodeEntry{d2ydx2}}, nullptr, &n));
+    return ret;
+});
 
 // abs
 MXNET_OPERATOR_REGISTER_UNARY_WITH_RSP_CSR(abs, cpu, mshadow_op::abs)
@@ -736,7 +767,26 @@ The storage type of ``abs`` output depends upon the input storage type:
 )code" ADD_FILELINE)
 .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_abs"});
 
-MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU(_backward_abs, unary_bwd<mshadow_op::sign>);
+MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU(_backward_abs, unary_bwd<mshadow_op::sign>)
+.set_attr<nnvm::FGradient>("FGradient",
+    [](const nnvm::NodePtr& n, const std::vector<nnvm::NodeEntry>& ograds) {
+      // ograds[0]: dL/dxgrad
+      // inputs[0]: dL/dy
+      // inputs[1]: x
+      // f(x) -> abs(x)
+      // f'(x) = 1 if x > 0 else -1
+      // f''(x) = 0
+      auto dydx = MakeNode("elemwise_div", n->attrs.name + "_dydx",
+                           {nnvm::NodeEntry{n}, n->inputs[0]}, nullptr, &n);
+
+      std::vector<nnvm::NodeEntry> ret;
+      ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad",
+                                {ograds[0], nnvm::NodeEntry(dydx)}, nullptr, &n));
+      ret.emplace_back(MakeNode("zeros_like", n->attrs.name + "_backward_grad_grad_in",
+                                {n->inputs[1]}, nullptr, &n));
+      return ret;
+    });
+
 
 // sign
 MXNET_OPERATOR_REGISTER_UNARY_WITH_RSP_CSR(sign, cpu, mshadow_op::sign)
diff --git a/tests/python/unittest/test_higher_order_grad.py b/tests/python/unittest/test_higher_order_grad.py
index ad14c50..0f07d01 100644
--- a/tests/python/unittest/test_higher_order_grad.py
+++ b/tests/python/unittest/test_higher_order_grad.py
@@ -107,6 +107,33 @@ def test_log10():
 
 
 @with_seed()
+def test_reciprocal():
+    def reciprocal(x):
+        return nd.reciprocal(x)
+
+    def grad_grad_op(x):
+        return 2 / x**3
+
+    for dim in range(1, 5):
+        shape = rand_shape_nd(dim)
+        array = random_arrays(shape)
+        check_second_order_unary(array, reciprocal, grad_grad_op)
+
+
+@with_seed()
+def test_abs():
+    def abs(x):
+        return nd.abs(x)
+
+    def grad_grad_op(x):
+        return nd.zeros_like(x)
+
+    for dim in range(1, 5):
+        shape = rand_shape_nd(dim)
+        array = random_arrays(shape)
+        check_second_order_unary(array, abs, grad_grad_op)
+
+
 def test_sigmoid():
     def sigmoid(x):
         return nd.sigmoid(x)