[Prim] Add stack_double_grad (#63161)

HydrogenSulfate · web-flow · commit 92f49a603aed · 2024-04-02T17:39:21.000+08:00
* add stack_double_grad composite API

* add TestStackDoubleGradCheck
diff --git a/paddle/fluid/eager/auto_code_generator/generator/codegen_utils.py b/paddle/fluid/eager/auto_code_generator/generator/codegen_utils.py
@@ -57,6 +57,7 @@
     "conv3d_double_grad",
     "depthwise_conv2d_grad_grad",
     "concat_double_grad",
+    "stack_double_grad",
     "expand_grad",
     "argsort_grad",
     "eigh_grad",
diff --git a/paddle/phi/api/yaml/backward.yaml b/paddle/phi/api/yaml/backward.yaml
@@ -2375,6 +2375,12 @@
   inplace : (out_grad -> x_grad)
   backward: squeeze_double_grad
 
+- backward_op : stack_double_grad
+  forward : stack_grad (Tensor[] x, Tensor grad_out, int axis=0) -> Tensor[](grad_x)
+  args : (Tensor[] grad_x_grad, int axis = 0)
+  output : Tensor(grad_out_grad)
+  invoke : stack(grad_x_grad, axis)
+
 - backward_op : stack_grad
   forward : stack (Tensor[] x, int axis) -> Tensor(out)
   args : (Tensor[] x, Tensor out_grad, int axis)
@@ -2389,6 +2395,7 @@
     data_type : out_grad
   no_need_buffer : x
   composite : stack_grad(x, out_grad, axis, x_grad)
+  backward: stack_double_grad
 
 - backward_op : stanh_grad
   forward : stanh(Tensor x, float scale_a, float scale_b) -> Tensor(out)
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -3225,7 +3225,7 @@
     outputs : [xshape]
 
 - op : stack
-  backward : stack_grad
+  backward : stack_grad, stack_double_grad
   inputs :
     x : X
   outputs :
diff --git a/test/legacy_test/test_nn_grad.py b/test/legacy_test/test_nn_grad.py
@@ -405,7 +405,6 @@ def concat_wrapper(self, x):
     @prog_scope()
     def func(self, place):
         x_shape = [2, 3, 4, 5]
-        pad = [1, 1, 1, 1]
         dtype = np.float64
 
         x1 = paddle.static.data('x', x_shape, dtype)
@@ -437,6 +436,45 @@ def test_grad(self):
             self.func(p)
 
 
+class TestStackDoubleGradCheck(unittest.TestCase):
+    def stack_wrapper(self, x):
+        return paddle.stack(x, axis=1)
+
+    @test_with_pir_api
+    @prog_scope()
+    def func(self, place):
+        x_shape = [2, 3, 4, 5]
+        dtype = np.float64
+
+        x1 = paddle.static.data('x', x_shape, dtype)
+        x2 = paddle.static.data('x', x_shape, dtype)
+        x1.persistable = True
+        x1.stop_gradient = False
+        x2.persistable = True
+        x2.stop_gradient = False
+        out = paddle.stack([x1, x2], axis=0)
+        x2_arr = np.random.uniform(-1, 1, x_shape).astype(dtype)
+        x1_arr = np.random.uniform(-1, 1, x_shape).astype(dtype)
+
+        gradient_checker.double_grad_check(
+            [x1, x2], out, x_init=[x1_arr, x2_arr], place=place
+        )
+        gradient_checker.double_grad_check_for_dygraph(
+            self.stack_wrapper,
+            [x1, x2],
+            out,
+            x_init=[x1_arr, x2_arr],
+            place=place,
+        )
+
+    def test_grad(self):
+        places = [base.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(base.CUDAPlace(0))
+        for p in places:
+            self.func(p)
+
+
 class TestAvgPool2DDoubleGradCheckCase1(unittest.TestCase):
     @test_with_pir_api
     @prog_scope()
