add acos_double_grad in dygraph composite (#70409)

zhiminzhang0830 · web-flow · commit fe4d0555090d · 2024-12-27T13:27:45.000+08:00
* add acos_double_grad in dygraph composite

* del acos at paddle/fluid/prim/api/api.yaml

* fix: update acos_double_grad

* fix: test_acos_double_grad_comp_dygraph

* fix bugs

* fix: chang infer_meta param
diff --git a/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py b/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
@@ -85,6 +85,7 @@
     "reshape_double_grad",
     "take_along_axis_double_grad",
     "index_add_double_grad",
+    "acos_double_grad",
 ]
 
 # white ops list whose kernel can automatically do type promotion.
diff --git a/paddle/fluid/pir/dialect/op_generator/vjp_interface_black_list.py b/paddle/fluid/pir/dialect/op_generator/vjp_interface_black_list.py
@@ -37,4 +37,5 @@
     'gather_nd_grad',
     'take_along_axis_grad',
     'index_add_grad',
+    'acos_grad',
 ]
diff --git a/paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h b/paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h
@@ -89,6 +89,26 @@ void cos_double_grad(const Tensor& x,
   }
 }
 
+template <typename T>
+void acos_double_grad(const Tensor& x,
+                      const Tensor& grad_out,
+                      const Tensor& grad_x_grad,
+                      Tensor* x_grad,
+                      Tensor* grad_out_grad) {
+  // acos grad grad : ddout = -((1-x*x)^(-0.5)) * ddx
+  // dx = dy * (-x)*((1-x*x)^(-1.5)) * ddx
+  auto x_tmp = 1 - x * x;
+  if (x_grad) {
+    auto x_grad_tmp = (grad_out * (-x) * pow<T>(x_tmp, -1.5) * grad_x_grad);
+    set_output<T>(x_grad_tmp, x_grad);
+  }
+
+  if (grad_out_grad) {
+    auto grad_out_grad_tmp = -pow<T>(x_tmp, -0.5) * grad_x_grad;
+    set_output<T>(grad_out_grad_tmp, grad_out_grad);
+  }
+}
+
 template <typename T>
 void minimum_double_grad(const Tensor& x,
                          const Tensor& y,
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -39,6 +39,18 @@
     support_trans_dtype : x
   composite : abs_triple_grad(x, grad_out_grad_grad, grad_x_grad_grad)
 
+- backward_op : acos_double_grad
+  forward : acos_grad (Tensor x, Tensor grad_out) -> Tensor(grad_x)
+  args : (Tensor x, Tensor grad_out, Tensor grad_x_grad)
+  output : Tensor(x_grad), Tensor(grad_out_grad)
+  infer_meta :
+    func : GeneralBinaryGradInferMeta
+    param : [x, grad_out]
+  kernel :
+    func : acos_double_grad
+  inplace : (grad_x_grad -> grad_out_grad)
+  composite : acos_double_grad(x, grad_out, grad_x_grad, x_grad, grad_out_grad)
+
 - backward_op : acos_grad
   forward : acos (Tensor x) -> Tensor(out)
   args : (Tensor x, Tensor out_grad)
@@ -49,6 +61,7 @@
   kernel :
     func : acos_grad
   inplace : (out_grad -> x_grad)
+  backward : acos_double_grad
 
 - backward_op : acosh_grad
   forward : acosh (Tensor x) -> Tensor(out)
diff --git a/test/prim/prim/vjp/eager/test_comp_eager_acos_double_grad.py b/test/prim/prim/vjp/eager/test_comp_eager_acos_double_grad.py
@@ -0,0 +1,54 @@
+# 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
+#
+#     https://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.uniform(-0.95, 0.95, size=(10, 10)),
+            np.random.rand(10, 10),
+            np.float32,
+        ),
+    ],
+)
+class TestAcosDoubleGradComp(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_acos_double_grad_comp_dygraph(self):
+        paddle.disable_static()
+        core.set_prim_eager_enabled(True)
+
+        x = paddle.to_tensor(self.primal, dtype='float32', stop_gradient=False)
+        x.stop_gradient = False
+        y = paddle.acos(x)
+        dx = paddle.grad(y, x, create_graph=True, retain_graph=True)[0]
+        ddx = paddle.grad(dx, x, create_graph=True, retain_graph=True)[0]
+        core.set_prim_eager_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -85,6 +85,7 @@`
`85`	`85`	`"reshape_double_grad",`
`86`	`86`	`"take_along_axis_double_grad",`
`87`	`87`	`"index_add_double_grad",`
	`88`	`+ "acos_double_grad",`
`88`	`89`	`]`
`89`	`90`
`90`	`91`	`# white ops list whose kernel can automatically do type promotion.`
Original file line number	Diff line number	Diff line change
`@@ -37,4 +37,5 @@`
`37`	`37`	`'gather_nd_grad',`
`38`	`38`	`'take_along_axis_grad',`
`39`	`39`	`'index_add_grad',`
	`40`	`+ 'acos_grad',`
`40`	`41`	`]`