support dynamic shape for expand_grad (#67481)

zeroRains · web-flow · commit 38a6189662c9 · 2024-08-19T09:20:15.000+08:00
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -1002,15 +1002,20 @@ void expand_grad(const Tensor& x,
                  const IntArray& shape,
                  Tensor* x_grad) {
   if (x_grad) {
-    if (out_grad.dims() != x.dims()) {
-      auto axes = get_reduce_dims_from_out(out_grad.dims(), x.dims());
-      auto reduced = out_grad.sum(common::vectorize(axes), x.dtype(), false);
-      if (reduced.dims().size() != x.dims().size()) {
-        reduced = reshape<T>(reduced, x.shape());
-      }
+    if (has_dynamic_shape(x.shape()) || has_dynamic_shape(out_grad.shape())) {
+      auto reduced = reduce_as<T>(out_grad, x);
       set_output<T>(reduced, x_grad);
     } else {
-      by_pass<T>(out_grad, x_grad);
+      if (out_grad.dims() != x.dims()) {
+        auto axes = get_reduce_dims_from_out(out_grad.dims(), x.dims());
+        auto reduced = out_grad.sum(common::vectorize(axes), x.dtype(), false);
+        if (reduced.dims().size() != x.dims().size()) {
+          reduced = reshape<T>(reduced, x.shape());
+        }
+        set_output<T>(reduced, x_grad);
+      } else {
+        by_pass<T>(out_grad, x_grad);
+      }
     }
   }
 }
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.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
@@ -199,6 +199,10 @@ def max_net6(x):
     return paddle.max(x)
 
 
+def expand_net(x):
+    return paddle.expand(x, [30, 200, 40])
+
+
 def apply_to_static(net, use_cinn, input_spec=None):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
@@ -2074,5 +2078,53 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimExpandWithGrad1(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [200, 40]
+        self.init_x_shape = [None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = expand_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimExpandWithGrad2(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 1, 40]
+        self.init_x_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = expand_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimExpandWithGrad3(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 1]
+        self.init_x_shape = [None, None, 1]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = expand_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimExpandWithGrad4(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 1, 1]
+        self.init_x_shape = [None, None, 1]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = expand_net
+        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.expand",`
`69`	`70`	`]`
`70`	`71`
`71`	`72`