support the dynamic shape for cumprod_grad op (#67965)

zeroRains · web-flow · commit f74a03b5616a · 2024-09-06T16:28:18.000+08:00
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -71,30 +71,37 @@ void cumprod_grad(const Tensor& x,
                   Tensor* x_grad) {
   if (x_grad) {
     // dx = cumsum(out * out_grad, dim, false, exclusive, !reverse) / x
-    std::vector<int64_t> x_dim = common::vectorize<int64_t>(x.dims());
-    auto zero_tensor = full<T>(x_dim, 0.0, x.dtype());
+    Tensor zero_tensor, ones_tensor;
+    if (has_dynamic_shape(x.shape())) {
+      zero_tensor = backend::full_with_tensor<T>(shape<T>(x), 0.0, x.dtype());
+      ones_tensor = backend::full_with_tensor<T>(shape<T>(x), 1.0, x.dtype());
+    } else {
+      zero_tensor = full<T>(x.shape(), 0.0, x.dtype());
+      ones_tensor = full<T>(x.shape(), 1.0, x.dtype());
+    }
     auto zero_mask = cast<T>(equal<T>(x, zero_tensor), x.dtype());
     // determine the index of first zero
     auto zero_mask_cumsum_exclusive =
         cumsum<T>(zero_mask, dim, false, true, reverse);
     auto zero_mask_cumsum = scale<T>(zero_mask_cumsum_exclusive, 2) + zero_mask;
-    auto ones_tensor = full<T>(x_dim, 1.0, x.dtype());
+
     auto first_zero_mask =
         cast<T>(equal<T>(zero_mask_cumsum, ones_tensor), x.dtype());
     // compute the grad for position with value not equal to 0
     auto common_dx = cumsum<T>(out * out_grad, dim, false, exclusive, !reverse);
     // fill the positions of 0 with 1.
-    auto replace_one = (1 - zero_mask) * x + zero_mask;
+    auto replace_one = (ones_tensor - zero_mask) * x + zero_mask;
     // fill the first positions of 0 with 1.
-    auto replace_first_one = (1 - first_zero_mask) * x + first_zero_mask;
+    auto replace_first_one =
+        (ones_tensor - first_zero_mask) * x + first_zero_mask;
     // recompute the grad of the first position with 0
     auto cumprod_recompute =
         cumprod<T>(replace_first_one, dim, exclusive, reverse);
     auto zeros_dx = cumsum<T>(
         cumprod_recompute * out_grad, dim, false, exclusive, !reverse);
-    auto x_grad_res =
-        ((1 - first_zero_mask) * common_dx + first_zero_mask * zeros_dx) /
-        replace_one;
+    auto x_grad_res = ((ones_tensor - first_zero_mask) * common_dx +
+                       first_zero_mask * zeros_dx) /
+                      replace_one;
     set_output<T>(x_grad_res, x_grad);
   }
 }
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -36,6 +36,7 @@
     "pd_op.cast",
     "pd_op.concat",
     "pd_op.cos",
+    "pd_op.cumprod",
     "pd_op.cumsum",
     "pd_op.divide",
     "pd_op.dropout",
diff --git a/test/prim/pir_prim/test_prim_sub_graph_abcde_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_abcde_backward_dynamic_shape.py
@@ -75,6 +75,10 @@ def concat_net3(x):
     return paddle.concat(x, axis=0)
 
 
+def cumprod_net(x):
+    return paddle.cumprod(x, 1)
+
+
 def cumsum_net(x):
     return paddle.cumsum(x, axis=1)
 
@@ -501,6 +505,18 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimCumprodWithGrad(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        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 = cumprod_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 class TestPrimCumsumWithGrad(TestPrimBaseWithGrad):
     def setUp(self):
         np.random.seed(2024)
