update square

Author: Eddie-Wang1120

paddle/fluid/primitive/codegen/gen.py

@@ -67,6 +67,7 @@
     'sin_grad',
     'cos_grad',
     'tanh_grad',
+    'square_grad',
 ]
 
 # prim op with two inputs and one output, with no attribute

paddle/fluid/primitive/rule/vjp/details.h

@@ -795,6 +795,14 @@ void log_grad(const Tensor& x, const Tensor& out_grad, Tensor* x_grad) {
   }
 }
 
+template <typename T>
+void square_grad(const Tensor& x, const Tensor& out_grad, Tensor* x_grad) {
+  if (x_grad) {
+    auto x_grad_tmp = 2 * x * out_grad;
+    set_output<T>(x_grad_tmp, x_grad);
+  }
+}
+
 template <typename T>
 void exp_grad(const Tensor& out, const Tensor& out_grad, Tensor* x_grad) {
   if (x_grad) {

test/legacy_test/test_activation_op.py

@@ -4297,6 +4297,7 @@ def test_check_grad(self):
             'Out',
             max_relative_error=0.007,
             check_pir=True,
+            check_prim_pir=True,
             check_pir_onednn=self.check_pir_onednn,
         )
 
@@ -4315,6 +4316,17 @@ def init_dtype(self):
     def test_check_output(self):
         self.check_output(check_pir=True)
 
+    def test_check_grad(self):
+        if self.dtype == np.float16:
+            return
+        self.check_grad(
+            ['X'],
+            'Out',
+            max_relative_error=0.007,
+            check_pir=True,
+            check_pir_onednn=self.check_pir_onednn,
+        )
+
 
 class TestSquare_Complex128(TestSquare):
     def init_dtype(self):
@@ -4323,6 +4335,17 @@ def init_dtype(self):
     def test_check_output(self):
         self.check_output(check_pir=True)
 
+    def test_check_grad(self):
+        if self.dtype == np.float16:
+            return
+        self.check_grad(
+            ['X'],
+            'Out',
+            max_relative_error=0.007,
+            check_pir=True,
+            check_pir_onednn=self.check_pir_onednn,
+        )
+
 
 class TestSquare_ZeroDim(TestSquare):
     def init_shape(self):
@@ -4365,7 +4388,12 @@ def test_check_output(self):
     def test_check_grad(self):
         place = core.CUDAPlace(0)
         self.check_grad_with_place(
-            place, ['X'], 'Out', numeric_grad_delta=0.5, check_pir=True
+            place,
+            ['X'],
+            'Out',
+            numeric_grad_delta=0.5,
+            check_pir=True,
+            check_prim_pir=True,
         )