Implement the composition of pow_double_grad (#62338)

YibinLiu666 · web-flow · commit f5a609c533f3 · 2024-03-26T14:25:01.000+08:00
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
@@ -114,6 +114,27 @@ void minimum_double_grad(const Tensor& x,
     }
   }
 }
+template <typename T>
+void pow_double_grad(const Tensor& x,
+                     const Tensor& grad_out,
+                     const Tensor& grad_x_grad,
+                     const Scalar& y,
+                     Tensor* x_grad,
+                     Tensor* grad_out_grad) {
+  // pow grad grad : ddout = y * pow(x, y-1) * ddx, dx = y * (y-1) * pow(x, y-2)
+  // * dout * ddx
+  auto y_value = y.to<float>();
+  if (grad_out_grad) {
+    auto grad_out_grad_tmp = y_value * x.pow(y_value - 1) * grad_x_grad;
+    set_output<T>(grad_out_grad_tmp, grad_out_grad);
+  }
+
+  if (x_grad) {
+    auto x_grad_tmp =
+        y_value * (y_value - 1) * x.pow(y_value - 2) * grad_out * grad_x_grad;
+    set_output<T>(x_grad_tmp, x_grad);
+  }
+}
 
 template <typename T>
 void maximum_double_grad(const Tensor& x,
diff --git a/paddle/phi/api/yaml/backward.yaml b/paddle/phi/api/yaml/backward.yaml
@@ -1772,6 +1772,7 @@
     data_type : x
   backward : pow_triple_grad
   inplace : (grad_x_grad -> x_grad)
+  composite: pow_double_grad(x, grad_out, grad_x_grad, y, x_grad, grad_out_grad)
 
 - backward_op : pow_grad
   forward : pow(Tensor x, Scalar y=1.0f) -> Tensor(out)
diff --git a/test/prim/prim/vjp/eager/test_comp_eager_pow_grad.py b/test/prim/prim/vjp/eager/test_comp_eager_pow_grad.py
@@ -18,6 +18,7 @@
 import unittest
 
 import numpy as np
+import parameterized as param
 from op_test import OpTest, convert_float_to_uint16
 
 import paddle
@@ -80,5 +81,51 @@ def if_enable_cinn(self):
         pass
 
 
+@param.parameterized_class(
+    ('primal', 'cotangent', 'dtype'),
+    [
+        (np.random.rand(10, 10), np.random.rand(10, 10), np.float32),
+    ],
+)
+class TestPowDoubleGradComp(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("pow_grad")
+            x = paddle.to_tensor(primal, dtype='float32', stop_gradient=False)
+            x.stop_gradient = False
+            y = paddle.pow(x, 2.7)
+            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.pow(x, 2.7)
+            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()
