Support the dynamic shape for pad_grad (#67606)

zeroRains · web-flow · commit 3a1ff8b7ccf8 · 2024-08-26T09:27:46.000+08:00
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -1552,21 +1552,40 @@ void pad_grad(const Tensor& input,
               const Scalar& pad_value,
               Tensor* input_grad) {
   if (input_grad) {
+    Tensor out_tmp;
     size_t rank = input.dims().size();
-    auto out_dims = out_grad.dims();
-
-    std::vector<int64_t> starts(rank, 0);
-    std::vector<int64_t> ends(rank, 0);
     std::vector<int64_t> axes(rank, 0);
     std::vector<int64_t> infer_flags(rank, 1);
     std::vector<int64_t> decrease_axis({});
-    for (size_t i = 0; i < rank; ++i) {
-      starts[i] = static_cast<int64_t>(paddings[2 * i]);
-      ends[i] = static_cast<int64_t>(out_dims[i] - paddings[2 * i + 1]);
-      axes[i] = i;
+    if (has_dynamic_shape(out_grad.shape())) {
+      auto out_shape = shape<T>(out_grad);
+      std::vector<Tensor> starts, ends;
+      for (size_t i = 0; i < rank; ++i) {
+        starts.push_back(full<T>({1}, paddings[2 * i], out_shape.dtype()));
+        ends.push_back(get_slice<T>(out_shape, i) -
+                       full<T>({1}, paddings[2 * i + 1], out_shape.dtype()));
+        axes[i] = i;
+      }
+      out_tmp = backend::slice<T>(out_grad,
+                                  concat<T>(starts),
+                                  concat<T>(ends),
+                                  axes,
+                                  infer_flags,
+                                  decrease_axis);
+    } else {
+      auto out_dims = out_grad.dims();
+
+      std::vector<int64_t> starts(rank, 0);
+      std::vector<int64_t> ends(rank, 0);
+
+      for (size_t i = 0; i < rank; ++i) {
+        starts[i] = static_cast<int64_t>(paddings[2 * i]);
+        ends[i] = static_cast<int64_t>(out_dims[i] - paddings[2 * i + 1]);
+        axes[i] = i;
+      }
+      out_tmp =
+          slice<T>(out_grad, axes, starts, ends, infer_flags, decrease_axis);
     }
-    auto out_tmp =
-        slice<T>(out_grad, axes, starts, ends, infer_flags, decrease_axis);
     set_output<T>(out_tmp, input_grad);
   }
 }
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -68,6 +68,7 @@
     "pd_op.max",
     "pd_op.stack",
     "pd_op.expand",
+    "pd_op.pad",
 ]
 
 
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
@@ -217,6 +217,10 @@ def stack_net3(x):
     return paddle.stack(x, axis=0)
 
 
+def pad_net(x):
+    return paddle.nn.functional.pad(x, [0, 1, 2])
+
+
 def apply_to_static(net, use_cinn, input_spec=None):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
@@ -2214,5 +2218,17 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimPadWithGrad(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [10, 20, 30, 40]
+        self.init_x_shape = [None, None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = pad_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -68,6 +68,7 @@`
`68`	`68`	`"pd_op.max",`
`69`	`69`	`"pd_op.stack",`
`70`	`70`	`"pd_op.expand",`
	`71`	`+ "pd_op.pad",`
`71`	`72`	`]`
`72`	`73`
`73`	`74`