surpport the dynamic shape for stack_grad (#67502)

zeroRains · web-flow · commit a996df5cec3e · 2024-08-20T10:43:25.000+08:00
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -830,20 +830,37 @@ void stack_grad(const std::vector<Tensor>& x,
   std::vector<int> sections(x_num, 1);
   std::vector<Tensor> x_grad_tmp =
       split<T>(out_grad, phi::IntArray(sections), axis);
-
-  // compose shape for each input tensor
-  std::vector<int64_t> grad_shape;
   auto out_dim = out_grad.dims().size();
-  for (int i = 0; i < out_dim; ++i) {
-    if (i != axis) {
-      grad_shape.push_back(out_grad.dims()[i]);
+  if (has_dynamic_shape(out_grad.shape())) {
+    Tensor out_grad_shape = shape<T>(out_grad);
+    std::vector<Tensor> grad_shape;
+    for (int i = 0; i < out_dim; i++) {
+      if (i != axis) {
+        grad_shape.push_back(get_slice<T>(out_grad_shape, i));
+      }
     }
-  }
+    Tensor grad_shape_tensor = concat<T>(grad_shape);
 
-  // assign to each input tensor if need grad(stop_gradient=False)
-  for (int i = 0; i < x_num; ++i) {
-    if (x_grad[i]) {
-      set_output<T>(reshape<T>(x_grad_tmp[i], grad_shape), x_grad[i]);
+    for (int i = 0; i < x_num; i++) {
+      if (x_grad[i]) {
+        set_output<T>(backend::reshape<T>(x_grad_tmp[i], grad_shape_tensor),
+                      x_grad[i]);
+      }
+    }
+  } else {
+    // compose shape for each input tensor
+    std::vector<int64_t> grad_shape;
+    for (int i = 0; i < out_dim; ++i) {
+      if (i != axis) {
+        grad_shape.push_back(out_grad.dims()[i]);
+      }
+    }
+
+    // assign to each input tensor if need grad(stop_gradient=False)
+    for (int i = 0; i < x_num; ++i) {
+      if (x_grad[i]) {
+        set_output<T>(reshape<T>(x_grad_tmp[i], grad_shape), x_grad[i]);
+      }
     }
   }
 }
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -66,6 +66,7 @@
     "pd_op.hardswish",
     "pd_op.reduce_as",
     "pd_op.max",
+    "pd_op.stack",
     "pd_op.expand",
 ]
 
diff --git a/test/prim/pir_prim/test_prim_sub_graph_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_backward_dynamic_shape.py
@@ -203,6 +203,20 @@ def expand_net(x):
     return paddle.expand(x, [30, 200, 40])
 
 
+def stack_net1(x):
+    y = x + 1
+    return paddle.stack([x, y], axis=-1)
+
+
+def stack_net2(x):
+    y = x + 1
+    return paddle.stack([x, y], axis=1)
+
+
+def stack_net3(x):
+    return paddle.stack(x, axis=0)
+
+
 def apply_to_static(net, use_cinn, input_spec=None):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
@@ -2126,5 +2140,79 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimStackWithGrad1(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = stack_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimStackWithGrad2(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = stack_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimStackWithGrad3(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = stack_net2
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimStackWithGrad4(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, 200, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = stack_net2
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimStackWithGrad5(TestPrimConcatWithGrad5):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        x = np.random.random(self.x_shape).astype(self.dtype)
+        self.x = [x + i for i in range(4)]
+        self.net = stack_net3
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimStackWithGrad6(TestPrimConcatWithGrad5):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, 200, None]
+        x = np.random.random(self.x_shape).astype(self.dtype)
+        self.x = [x + i for i in range(4)]
+        self.net = stack_net3
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -66,6 +66,7 @@`
`66`	`66`	`"pd_op.hardswish",`
`67`	`67`	`"pd_op.reduce_as",`
`68`	`68`	`"pd_op.max",`
	`69`	`+ "pd_op.stack",`
`69`	`70`	`"pd_op.expand",`
`70`	`71`	`]`
`71`	`72`