[MXNET-978] Second order gradient support for some unary operators

src/operator/tensor/elemwise_unary_op_basic.cc

@@ -85,8 +85,20 @@ The storage type of ``relu`` output depends upon the input storage type:
 )code" ADD_FILELINE)
 .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseOut{"_backward_relu"});
 
-MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU(_backward_relu,
-                                               unary_bwd<mshadow_op::relu_grad>);
+MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU(_backward_relu, unary_bwd<mshadow_op::relu_grad>)
+.set_attr<nnvm::FGradient>("FGradient",
+    [](const nnvm::NodePtr& n, const std::vector<nnvm::NodeEntry>& ograds) {
+      std::vector<nnvm::NodeEntry> ret;
+      // f(x) -> f = relu
+      // f'(x) = 1 if x > 0 else 0
+      // f''(x) = 0
+      auto gx = nnvm::NodeEntry{n};  // f'(x)
+      ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad",
+                                {ograds[0], gx}, nullptr, &n));
+      ret.emplace_back(MakeNode("zeros_like", n->attrs.name + "_backward_grad_grad_in",
+                                {gx}, nullptr, &n));
+      return ret;
+    });
 
 // sigmoid
 MXNET_OPERATOR_REGISTER_UNARY(sigmoid)

src/operator/tensor/elemwise_unary_op_trig.cc

@@ -46,7 +46,26 @@ The storage type of ``sin`` output depends upon the input storage type:
 )code" ADD_FILELINE)
 .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{ "_backward_sin" });
 
-MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU_DR(_backward_sin, unary_bwd<mshadow_op::sin_grad>);
+MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU_DR(_backward_sin, unary_bwd<mshadow_op::sin_grad>)
+.set_attr<nnvm::FGradient>("FGradient",
+    [](const nnvm::NodePtr& n, const std::vector<nnvm::NodeEntry>& ograds) {
+      // f(x) = sin(x)
+      // f'(x) = cos(x)
+      // f''(x) = -sin(x)
+      auto grad_x = nnvm::NodeEntry(n);
+      auto grad_grad_x_mid = MakeNode("sin", n->attrs.name + "_mid_grad_grad",
+                                      {n->inputs[1]}, nullptr, &n);
+      auto grad_grad_x = MakeNode("negative", n->attrs.name + "_backward_grad_grad",
+                                  {nnvm::NodeEntry(grad_grad_x_mid)}, nullptr, &n);
+      std::vector<nnvm::NodeEntry> ret;
+      // for the backward of the _backward_sin node
+      // first input is the ograd and second input is x (because ElemwiseUseIn)
+      ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad",
+                                {ograds[0], grad_x}, nullptr, &n));
+      ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad_in",
+                                {ograds[0], nnvm::NodeEntry(grad_grad_x)}, nullptr, &n));
+      return ret;
+    });
 
 // cos
 MXNET_OPERATOR_REGISTER_UNARY_WITH_SPARSE_DR(cos, cpu, mshadow_op::cos)
@@ -63,7 +82,27 @@ The storage type of ``cos`` output is always dense
 )code" ADD_FILELINE)
 .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_cos"});
 
-MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU(_backward_cos, unary_bwd<mshadow_op::cos_grad>);
+MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU(_backward_cos, unary_bwd<mshadow_op::cos_grad>)
+.set_attr<nnvm::FGradient>("FGradient",
+    [](const nnvm::NodePtr& n, const std::vector<nnvm::NodeEntry>& ograds) {
+      // f(x) = cos(x)
+      // f'(x) = -sin(x)
+      // f''(x) = -cos(x)
+      auto grad_x = nnvm::NodeEntry(n);
+      auto grad_grad_x_mid = MakeNode("cos", n->attrs.name + "_mid_grad_grad",
+                                      {n->inputs[1]}, nullptr, &n);
+      auto grad_grad_x = MakeNode("negative", n->attrs.name + "_backward_grad_grad",
+                                  {nnvm::NodeEntry(grad_grad_x_mid)}, nullptr, &n);
+      std::vector<nnvm::NodeEntry> ret;
+      // for the backward of the _backward_cos node
+      // first input is the ograd and second input is x (because ElemwiseUseIn)
+      ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad",
+                                {ograds[0], grad_x}, nullptr, &n));
+      ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad_in",
+                                {ograds[0], nnvm::NodeEntry(grad_grad_x)}, nullptr, &n));
+      return ret;
+    });
+
 
 // tan
 MXNET_OPERATOR_REGISTER_UNARY_WITH_RSP_CSR(tan, cpu, mshadow_op::tan)

tests/python/unittest/test_higher_order_grad.py

@@ -15,13 +15,52 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import math
 
+import math
 from mxnet import nd, autograd
 from mxnet.test_utils import assert_almost_equal, random_arrays
 from common import with_seed
 
 
+@with_seed()
+def test_sin():
+    def sin(x):
+        return nd.sin(x)
+
+    def grad_grad_op(x):
+        return -nd.sin(x)
+
+    arrays = random_arrays((2, 2), (2, 3), (4, 5, 2), (3, 1, 4, 5))
+    for array in arrays:
+        check_second_order_unary(array, sin, grad_grad_op)
+
+
+@with_seed()
+def test_cos():
+    def cos(x):
+        return nd.cos(x)
+
+    def grad_grad_op(x):
+        return -nd.cos(x)
+
+    arrays = random_arrays((2, 2), (2, 3), (4, 5, 2), (3, 1, 4, 5))
+    for array in arrays:
+        check_second_order_unary(array, cos, grad_grad_op)
+
+
+@with_seed()
+def test_relu():
+    def relu(x):
+        return nd.relu(x)
+
+    def grad_grad_op(x):
+        return nd.zeros_like(x)
+
+    arrays = random_arrays((2, 2), (2, 3), (4, 5, 2), (3, 1, 4, 5))
+    for array in arrays:
+        check_second_order_unary(array, relu, grad_grad_op)
+
+
 @with_seed()
 def test_log():
     def log(x):