[bug] fix higher grad log

src/operator/tensor/elemwise_unary_op_basic.cc

@@ -1074,16 +1074,19 @@ MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU_DR(_backward_log,
   [](const nnvm::NodePtr& n, const std::vector<nnvm::NodeEntry>& ograds) {
     // For f(x) -> f = log
     // f''(x) = -1 * (f'(x) * f'(x))
-    auto gx = nnvm::NodeEntry{n};
+    auto gx_mul_head_grads = nnvm::NodeEntry{n};  // f'(x) * head_grads
+    auto head_grads = nnvm::NodeEntry{n->inputs[0]};
+    auto g_lx = MakeNode("reciprocal", n->attrs.name + "_backward_log_grad",
+                            {n->inputs[1]}, nullptr, &n);
     auto ggx_mid = MakeNode("elemwise_mul", n->attrs.name + "_backward_mid_grad_grad",
-                            {gx, gx}, nullptr, &n);
+                            {gx_mul_head_grads, nnvm::NodeEntry{g_lx}}, nullptr, &n);
     auto ggx = MakeNode("negative", n->attrs.name + "_backward_grad_grad",
                         {nnvm::NodeEntry{ggx_mid}}, nullptr, &n);
 
     std::vector<nnvm::NodeEntry> ret;
 
     ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad",
-                             {ograds[0], gx}, nullptr, &n));
+                             {ograds[0], nnvm::NodeEntry{g_lx}}, nullptr, &n));
     ret.emplace_back(MakeNode("elemwise_mul", n->attrs.name + "_backward_grad_grad_inp",
                              {ograds[0], nnvm::NodeEntry{ggx}}, nullptr, &n));
     return ret;

tests/python/unittest/test_higher_order_grad.py

@@ -27,52 +27,79 @@ def test_log():
     def log(x):
         return nd.log(x)
 
+    def grad_op(x):
+        return 1/x
+
     def grad_grad_op(x):
         return -1/(x**2)
 
     arrays = random_arrays((2, 2), (2, 3), (4, 5, 2), (3, 1, 4, 5))
 
     for array in arrays:
-        check_second_order_unary(array, log, grad_grad_op)
+        check_second_order_unary(array, log, grad_op, grad_grad_op)
 
 
 @with_seed()
 def test_log2():
     def log2(x):
         return nd.log2(x)
 
+    def grad_op(x):
+        return 1/(x * math.log(2))
+
     def grad_grad_op(x):
         return -1/((x**2) * math.log(2))
 
     arrays = random_arrays((2, 2), (2, 3), (4, 5, 2), (3, 1, 4, 5))
 
     for array in arrays:
-        check_second_order_unary(array, log2, grad_grad_op)
+        check_second_order_unary(array, log2, grad_op, grad_grad_op)
 
 
 @with_seed()
 def test_log10():
     def log10(x):
         return nd.log10(x)
 
+    def grad_op(x):
+        return 1/(x * math.log(10))
+
     def grad_grad_op(x):
         return -1/((x**2) * math.log(10))
 
     arrays = random_arrays((2, 2), (2, 3), (4, 5, 2), (3, 1, 4, 5))
 
     for array in arrays:
-        check_second_order_unary(array, log10, grad_grad_op)
+        check_second_order_unary(array, log10, grad_op, grad_grad_op)
 
 
-def check_second_order_unary(x, op, grad_grad_op):
+def check_second_order_unary(x, op, grad_op, grad_grad_op):
     x = nd.array(x)
-    expect_grad_grad = grad_grad_op(x)
+    grad_x = grad_op(x)
+    grad_grad_x = grad_grad_op(x)
     x.attach_grad()
+
+    # Manual head_grads.
+    head_grads = nd.random.normal(shape=x.shape)
+    head_grad_grads = nd.random.normal(shape=x.shape)
+    head_grads.attach_grad()
+
+    # Perform compute.
     with autograd.record():
         y = op(x)
-        y_grad = autograd.grad(y, x, create_graph=True, retain_graph=True)[0]
-    y_grad.backward()
-    assert_almost_equal(expect_grad_grad.asnumpy(), x.grad.asnumpy())
+        y_grad = autograd.grad(y, x, head_grads=head_grads,
+                               create_graph=True, retain_graph=True)[0]
+
+    y_grad.backward(head_grad_grads)
+
+    # Compute expected values.
+    expected_grad_grad = grad_grad_x.asnumpy() * head_grad_grads.asnumpy() * \
+        head_grads.asnumpy()
+    expected_heads_grad = grad_x.asnumpy()
+
+    # Validate the gradients.
+    assert_almost_equal(expected_grad_grad, x.grad.asnumpy())
+    assert_almost_equal(expected_heads_grad, head_grads.grad.asnumpy())
 
 
 if __name__ == '__main__':