Dev sin loop grad (#7)

* edit tanh to a closure op

* add grad-looped sin_cos_negative

* add test case

Co-authored-by: yoonlee888 <qiuyunlei@zhejianglab.com>
Co-authored-by: Zhenhua <1209435+hengzi@users.noreply.github.com>

oneflow/core/autograd/gradient_funcs/cos.cpp

@@ -0,0 +1,53 @@
+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+#include "oneflow/core/framework/op_expr_grad_function.h"
+#include "oneflow/core/framework/op_interpreter/op_interpreter_util.h"
+#include "oneflow/core/functional/functional.h"
+
+namespace oneflow {
+namespace one {
+
+struct CosCaptureState : public AutoGradCaptureState {
+  bool x_requires_grad;
+};
+
+class Cos : public OpExprGradFunction<CosCaptureState> {
+  Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
+
+  Maybe<void> Capture(CosCaptureState* ctx, const TensorTuple& inputs,
+                      const TensorTuple& outputs, const AttrMap& attrs) const override {
+    ctx->x_requires_grad = inputs.at(0)->requires_grad();
+    ctx->SaveTensorForBackward(inputs.at(0));
+    return Maybe<void>::Ok();
+  }
+
+  Maybe<void> Apply(const CosCaptureState* ctx, const TensorTuple& out_grads,
+                    TensorTuple* in_grads) const override {
+    if (!ctx->x_requires_grad) { return Maybe<void>::Ok(); }
+
+    const auto& x = ctx->SavedTensors().at(0);
+    const auto& xx = JUST(functional::Sin(x));
+    const auto& xxx = JUST(functional::Negative(xx));
+    in_grads->at(0) = JUST(functional::Mul(out_grads.at(0), xxx));
+    return Maybe<void>::Ok();
+  }
+};
+
+REGISTER_OP_EXPR_GRAD_FUNCTION("cos", Cos);
+
+}  // namespace one
+}  // namespace oneflow

oneflow/core/autograd/gradient_funcs/math_unary_op.cpp

@@ -56,7 +56,7 @@ class UnaryMathOp : public OpExprGradFunction<UnaryMathCaptureState> {
   class op_cls##Cls final : public UnaryMathOp<functional::op_cls##Grad> {}; \
   REGISTER_OP_EXPR_GRAD_FUNCTION(op_type_name, op_cls##Cls);
 
-OF_PP_FOR_EACH_TUPLE(INSTANTIAT_AND_REGISTER_UNARY_MATHOP_CLASS, MATH_UNARY_ELEMENTWISE_FUNC_SEQ);
+OF_PP_FOR_EACH_TUPLE(INSTANTIAT_AND_REGISTER_UNARY_MATHOP_CLASS, MATH_UNARY_ELEMENTWISE_GRAD_SEQ);
 
 #undef INSTANTIAT_AND_REGISTER_UNARY_MATHOP_CLASS
 }  // namespace one

oneflow/core/autograd/gradient_funcs/negative.cpp

@@ -0,0 +1,49 @@
+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+#include "oneflow/core/framework/op_expr_grad_function.h"
+#include "oneflow/core/framework/op_interpreter/op_interpreter_util.h"
+#include "oneflow/core/functional/functional.h"
+
+namespace oneflow {
+namespace one {
+
+struct NegativeCaptureState : public AutoGradCaptureState {
+  bool x_requires_grad;
+};
+
+class Negative : public OpExprGradFunction<NegativeCaptureState> {
+  Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
+
+  Maybe<void> Capture(NegativeCaptureState* ctx, const TensorTuple& inputs,
+                      const TensorTuple& outputs, const AttrMap& attrs) const override {
+    ctx->x_requires_grad = inputs.at(0)->requires_grad();
+    ctx->SaveTensorForBackward(inputs.at(0));
+    return Maybe<void>::Ok();
+  }
+
+  Maybe<void> Apply(const NegativeCaptureState* ctx, const TensorTuple& out_grads,
+                    TensorTuple* in_grads) const override {
+    if (!ctx->x_requires_grad) { return Maybe<void>::Ok(); }
+    in_grads->at(0) = JUST(functional::Negative(out_grads.at(0)));
+    return Maybe<void>::Ok();
+  }
+};
+
+REGISTER_OP_EXPR_GRAD_FUNCTION("negative", Negative);
+
+}  // namespace one
+}  // namespace oneflow

oneflow/core/autograd/gradient_funcs/sin.cpp

@@ -0,0 +1,52 @@
+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+#include "oneflow/core/framework/op_expr_grad_function.h"
+#include "oneflow/core/framework/op_interpreter/op_interpreter_util.h"
+#include "oneflow/core/functional/functional.h"
+
+namespace oneflow {
+namespace one {
+
+struct SinCaptureState : public AutoGradCaptureState {
+  bool x_requires_grad;
+};
+
+class Sin : public OpExprGradFunction<SinCaptureState> {
+  Maybe<void> Init(const OpExpr& op) override { return Maybe<void>::Ok(); }
+
+  Maybe<void> Capture(SinCaptureState* ctx, const TensorTuple& inputs,
+                      const TensorTuple& outputs, const AttrMap& attrs) const override {
+    ctx->x_requires_grad = inputs.at(0)->requires_grad();
+    ctx->SaveTensorForBackward(inputs.at(0));
+    return Maybe<void>::Ok();
+  }
+
+  Maybe<void> Apply(const SinCaptureState* ctx, const TensorTuple& out_grads,
+                    TensorTuple* in_grads) const override {
+    if (!ctx->x_requires_grad) { return Maybe<void>::Ok(); }
+
+    const auto& x = ctx->SavedTensors().at(0);
+    const auto& xx = JUST(functional::Cos(x));
+    in_grads->at(0) = JUST(functional::Mul(out_grads.at(0), xx));
+    return Maybe<void>::Ok();
+  }
+};
+
+REGISTER_OP_EXPR_GRAD_FUNCTION("sin", Sin);
+
+}  // namespace one
+}  // namespace oneflow

oneflow/user/ops/math_unary_elementwise_seq.h

@@ -92,6 +92,38 @@ namespace oneflow {
   OF_PP_MAKE_TUPLE_SEQ("tan", Tan)                           \
   OF_PP_MAKE_TUPLE_SEQ("not_equal_zero", NotEqualZero)
 
+#define MATH_UNARY_ELEMENTWISE_GRAD_SEQ                      \
+  OF_PP_MAKE_TUPLE_SEQ("abs", Abs)                           \
+  OF_PP_MAKE_TUPLE_SEQ("acos", Acos)                         \
+  OF_PP_MAKE_TUPLE_SEQ("acosh", Acosh)                       \
+  OF_PP_MAKE_TUPLE_SEQ("asin", Asin)                         \
+  OF_PP_MAKE_TUPLE_SEQ("asinh", Asinh)                       \
+  OF_PP_MAKE_TUPLE_SEQ("atan", Atan)                         \
+  OF_PP_MAKE_TUPLE_SEQ("atanh", Atanh)                       \
+  OF_PP_MAKE_TUPLE_SEQ("ceil", Ceil)                         \
+  OF_PP_MAKE_TUPLE_SEQ("cosh", Cosh)                         \
+  OF_PP_MAKE_TUPLE_SEQ("erf", Erf)                           \
+  OF_PP_MAKE_TUPLE_SEQ("erfc", Erfc)                         \
+  OF_PP_MAKE_TUPLE_SEQ("exp", Exp)                           \
+  OF_PP_MAKE_TUPLE_SEQ("expm1", Expm1)                       \
+  OF_PP_MAKE_TUPLE_SEQ("floor", Floor)                       \
+  OF_PP_MAKE_TUPLE_SEQ("lgamma", Lgamma)                     \
+  OF_PP_MAKE_TUPLE_SEQ("log", Log)                           \
+  OF_PP_MAKE_TUPLE_SEQ("log2", Log2)                         \
+  OF_PP_MAKE_TUPLE_SEQ("log1p", Log1p)                       \
+  OF_PP_MAKE_TUPLE_SEQ("log_sigmoid", LogSigmoid)            \
+  OF_PP_MAKE_TUPLE_SEQ("reciprocal", Reciprocal)             \
+  OF_PP_MAKE_TUPLE_SEQ("reciprocal_no_nan", ReciprocalNoNan) \
+  OF_PP_MAKE_TUPLE_SEQ("rint", Rint)                         \
+  OF_PP_MAKE_TUPLE_SEQ("round", Round)                       \
+  OF_PP_MAKE_TUPLE_SEQ("rsqrt", Rsqrt)                       \
+  OF_PP_MAKE_TUPLE_SEQ("sigmoid_v2", Sigmoid)                \
+  OF_PP_MAKE_TUPLE_SEQ("sign", Sign)                         \
+  OF_PP_MAKE_TUPLE_SEQ("sinh", Sinh)                         \
+  OF_PP_MAKE_TUPLE_SEQ("sqrt", Sqrt)                         \
+  OF_PP_MAKE_TUPLE_SEQ("square", Square)                     \
+  OF_PP_MAKE_TUPLE_SEQ("tan", Tan)                           \
+  OF_PP_MAKE_TUPLE_SEQ("not_equal_zero", NotEqualZero)
 }  // namespace oneflow
 
 #endif  // ONEFLOW_USER_OPS_MATH_UNARY_ELEMENTWISE_SEQ_H_

python/oneflow/test/modules/test_math_ops.py

@@ -82,12 +82,38 @@ def _test_cos_backward(test_case, shape, device):
     np_grad = -np.sin(x.numpy())
     test_case.assertTrue(np.allclose(x.grad.numpy(), np_grad, 1e-05, 1e-05))
 
+def _test_cos_higher_order_backward(test_case, shape, device):
+    x = flow.tensor(
+        np.random.randn(*shape),
+        dtype=flow.float32,
+        device=flow.device(device),
+        requires_grad=True,
+    )
+    y = flow.cos(x)
+    x_grad = flow.autograd.grad(
+        outputs=y,
+        inputs=x,
+        grad_outputs=flow.ones_like(y),
+        create_graph = True
+        )[0]
+    np_grad = -np.sin(x.numpy())
+    test_case.assertTrue(np.allclose(x_grad, np_grad, 1e-05, 1e-05))
+
+    x_grad_grad = flow.autograd.grad(
+        outputs=x_grad,
+        inputs=x,
+        grad_outputs=flow.ones_like(y),
+        create_graph = True
+        )[0]
+    np_grad_grad = -np.cos(x.numpy())
+    test_case.assertTrue(np.allclose(x_grad_grad, np_grad_grad, 1e-05, 1e-05))
+
 
 @flow.unittest.skip_unless_1n1d()
 class TestCos(flow.unittest.TestCase):
     def test_cos(test_case):
         arg_dict = OrderedDict()
-        arg_dict["test_fun"] = [_test_cos, _test_cos_backward]
+        arg_dict["test_fun"] = [_test_cos, _test_cos_backward, _test_cos_higher_order_backward]
         arg_dict["shape"] = [(2, 3), (2, 3, 4), (2, 3, 4, 5)]
         arg_dict["device"] = ["cpu", "cuda"]
         for arg in GenArgList(arg_dict):