Implement the composition of cos_double_grad

paddle/fluid/eager/auto_code_generator/generator/codegen_utils.py

@@ -39,7 +39,6 @@
     "tanh_double_grad",
     "tanh_triple_grad",
     "sin_triple_grad",
-    "cos_double_grad",
     "cos_triple_grad",
     "subtract_double_grad",
     "divide_double_grad",

paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h

@@ -71,6 +71,24 @@ void sin_double_grad(const Tensor& x,
   }
 }
 
+template <typename T>
+void cos_double_grad(const Tensor& x,
+                     const Tensor& grad_out,
+                     const Tensor& grad_x_grad,
+                     Tensor* x_grad,
+                     Tensor* grad_out_grad) {
+  // cos grad grad : ddout = -sinx * ddx, dx = -dy * cosx * ddx
+  if (x_grad) {
+    auto x_grad_tmp = -(grad_out * cos<T>(x) * grad_x_grad);
+    set_output<T>(x_grad_tmp, x_grad);
+  }
+
+  if (grad_out_grad) {
+    auto grad_out_grad_tmp = -sin<T>(x) * grad_x_grad;
+    set_output<T>(grad_out_grad_tmp, grad_out_grad);
+  }
+}
+
 template <typename T>
 void tanh_triple_grad(const Tensor& out,
                       const Tensor& grad_out_forward,

paddle/phi/api/yaml/backward.yaml

@@ -433,9 +433,9 @@
     param : [x, x]
   kernel :
     func : cos_double_grad
-  optional: grad_out
   backward : cos_triple_grad
   inplace : (grad_x_grad -> grad_out_grad)
+  composite : cos_double_grad(x, grad_out, grad_x_grad, x_grad, grad_out_grad)
 
 - backward_op : cos_grad
   forward : cos (Tensor x) -> Tensor(out)

paddle/phi/kernels/activation_grad_kernel.h

@@ -97,7 +97,7 @@ void SinDoubleGradKernel(const Context& dev_ctx,
 template <typename T, typename Context>
 void CosDoubleGradKernel(const Context& dev_ctx,
                          const DenseTensor& x,
-                         const paddle::optional<DenseTensor>& dout,
+                         const DenseTensor& dout,
                          const DenseTensor& ddx,
                          DenseTensor* dx,
                          DenseTensor* ddout);

paddle/phi/kernels/impl/activation_grad_impl.h

@@ -717,7 +717,7 @@ void SinTripleGradKernel(const Context& dev_ctx,
 template <typename T, typename Context>
 void CosDoubleGradKernel(const Context& dev_ctx,
                          const DenseTensor& x,
-                         const paddle::optional<DenseTensor>& dout,
+                         const DenseTensor& dout,
                          const DenseTensor& ddx,
                          DenseTensor* dx,
                          DenseTensor* ddout) {
@@ -728,7 +728,7 @@ void CosDoubleGradKernel(const Context& dev_ctx,
     dev_ctx.template Alloc<T>(ddout);
   }
   phi::funcs::CosDoubleGradFunctor<T> functor;
-  functor(dev_ctx, &x, dout.get_ptr(), &ddx, dx, ddout);
+  functor(dev_ctx, &x, &dout, &ddx, dx, ddout);
 }
 
 template <typename T, typename Context>

test/prim/prim/vjp/eager/test_comp_eager_cos_double_grad.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2024 PaddlePaddle 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.
+
+import unittest
+
+import numpy as np
+import parameterized as param
+
+import paddle
+from paddle.base import core
+
+
+@param.parameterized_class(
+    ('primal', 'cotangent', 'dtype'),
+    [
+        (np.random.rand(10, 10), np.random.rand(10, 10), np.float32),
+    ],
+)
+class TestCosDoubleGradComp(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.primal = cls.primal.astype(cls.dtype)
+        if cls.cotangent is not None:
+            cls.cotangent = cls.cotangent.astype(cls.dtype)
+
+    def test_cos_double_grad_comp_dygraph(self):
+        def actual(primal):
+            paddle.disable_static()
+            core.set_prim_eager_enabled(True)
+            core._set_prim_backward_blacklist("cos_grad")
+            x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y = paddle.cos(x)
+            dx = paddle.grad(y, x, create_graph=True, retain_graph=True)
+
+            ddx = paddle.grad(dx, x, create_graph=True, retain_graph=True)
+            return ddx[0]
+
+        def desired(primal):
+            paddle.disable_static()
+            core.set_prim_eager_enabled(False)
+            x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y = paddle.cos(x)
+            dx = paddle.grad(y, x, create_graph=True, retain_graph=True)
+
+            ddx = paddle.grad(dx, x, create_graph=True, retain_graph=True)
+            return ddx[0]
+
+        np.testing.assert_allclose(
+            actual=actual(self.primal),
+            desired=desired(self.primal),
+            rtol=1e-6,
+            atol=0,
+        )
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()