fix bugs

oneflow/core/autograd/gradient_funcs/tanh.cpp

@@ -29,8 +29,8 @@ struct TanhCaptureState : public AutoGradCaptureState {
 class TanhGrad : public OpExprGradFunction<TanhCaptureState> {
   Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
 
-  Maybe<void> Capture(TanhCaptureState* ctx, const TensorTuple& inputs,
-                      const TensorTuple& outputs, const AttrMap& attrs) const override {
+  Maybe<void> Capture(TanhCaptureState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
+                      const AttrMap& attrs) const override {
     ctx->x_requires_grad = inputs.at(0)->requires_grad();
     ctx->SaveTensorForBackward(outputs.at(0));
     return Maybe<void>::Ok();
@@ -40,9 +40,8 @@ class TanhGrad : public OpExprGradFunction<TanhCaptureState> {
                     TensorTuple* in_grads) const override {
     if (!ctx->x_requires_grad) { return Maybe<void>::Ok(); }
     const auto& y = ctx->SavedTensors().at(0);
-    const auto& a = functional::Mul(y, y);
-    const auto& aa = functional::ScalarSub(1, JUST(a));
-    in_grads->at(0) = JUST(functional::Mul(out_grads.at(0), JUST(aa)));
+    in_grads->at(0) = JUST(functional::Mul(
+        out_grads.at(0), JUST(functional::ScalarSub(1, JUST(functional::Mul(y, y)), 1))));
     return Maybe<void>::Ok();
   }
 };

oneflow/core/functional/impl/higher_derivative_functor.cpp

@@ -71,7 +71,7 @@ class ExpGradGradFunctor {
 class LogGradGradFunctor {
  public:
   Maybe<Tensor> operator()(const std::shared_ptr<Tensor>& x,
-                           const std::shared_ptr<Tensor>& dydx) const {                 
+                           const std::shared_ptr<Tensor>& dydx) const {
     auto res = sequence_function(functional::Square)
                    .then(functional::Negative)
                    .then(functional::Reciprocal)

python/oneflow/test/modules/test_consistent_math_op_higher_derivative.py

@@ -58,6 +58,24 @@ def test_consistent_cos_grad_grad(test_case):
             for sbp in all_sbp(placement, max_dim=2):
                 _consistent_math_op_grad_grad_impl(test_case, "cos", placement, sbp)
 
+    @globaltest
+    def test_consistent_exp_grad_grad(test_case):
+        for placement in all_placement():
+            for sbp in all_sbp(placement, max_dim=2):
+                _consistent_math_op_grad_grad_impl(test_case, "exp", placement, sbp)
+
+    @globaltest
+    def test_consistent_log_grad_grad(test_case):
+        for placement in all_placement():
+            for sbp in all_sbp(placement, max_dim=2):
+                _consistent_math_op_grad_grad_impl(test_case, "log", placement, sbp)
+
+    @globaltest
+    def test_consistent_tanh_grad_grad(test_case):
+        for placement in all_placement():
+            for sbp in all_sbp(placement, max_dim=2):
+                _consistent_math_op_grad_grad_impl(test_case, "tanh", placement, sbp)
+
 
 if __name__ == "__main__":
     unittest.main()

python/oneflow/test/modules/test_math_op_higher_derivative.py

@@ -50,9 +50,13 @@ def test_cos_grad_grad(test_case):
 
     def test_exp_grad_grad(test_case):
         _test_math_op_grad_grad_impl(test_case, "exp")
-    
+
     def test_log_grad_grad(test_case):
         _test_math_op_grad_grad_impl(test_case, "log")
 
+    def test_tanh_grad_grad(test_case):
+        _test_math_op_grad_grad_impl(test_case, "tanh")
+
+
 if __name__ == "__main__":
     unittest.main()