surport the dynamic shape for reduce_as (#67151)

zeroRains · web-flow · commit e29c92c86b8b · 2024-08-09T10:16:59.000+08:00
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -369,36 +369,41 @@ void reduce_as_grad(const Tensor& x,
   if (!x_grad) {
     return;
   }
-  std::vector<int64_t> x_dim = common::vectorize<int64_t>(x.dims());
-  std::vector<int64_t> axis = common::vectorize<int64_t>(
-      get_reduce_dims_from_out(x.dims(), target.dims()));
-  int64_t axis_size = axis.size();
-  if (axis_size == 0) {
-    by_pass<T>(out_grad, x_grad);
-    return;
-  }
-  int64_t x_dim_size = x_dim.size();
-
-  auto x_grad_tmp = Tensor();
-  if (x_dim_size == 1) {
-    x_grad_tmp = expand<T>(out_grad, IntArray(x_dim));
+  if (has_dynamic_shape(x.shape()) || has_dynamic_shape(out_grad.shape())) {
+    auto x_grad_tmp = backend::expand_with_tensor<T>(out_grad, shape<T>(x));
+    set_output<T>(x_grad_tmp, x_grad);
   } else {
-    auto axis_ = std::vector<int64_t>();
-    for (int64_t i = 0; i < axis_size; i++) {
-      axis_.push_back(axis[i]);
-      if (axis[i] < 0) {
-        axis_[i] += x_dim_size;
-      }
+    std::vector<int64_t> x_dim = common::vectorize<int64_t>(x.dims());
+    std::vector<int64_t> axis = common::vectorize<int64_t>(
+        get_reduce_dims_from_out(x.dims(), target.dims()));
+    int64_t axis_size = axis.size();
+    if (axis_size == 0) {
+      by_pass<T>(out_grad, x_grad);
+      return;
     }
-    Tensor out_grad_ = out_grad;
-    if (out_grad.shape().size() != x.shape().size()) {
-      auto out_grad_shape = get_unsqueeze_dims(out_grad, axis_);
-      out_grad_ = reshape<T>(out_grad, out_grad_shape);
+    int64_t x_dim_size = x_dim.size();
+
+    auto x_grad_tmp = Tensor();
+    if (x_dim_size == 1) {
+      x_grad_tmp = expand<T>(out_grad, IntArray(x_dim));
+    } else {
+      auto axis_ = std::vector<int64_t>();
+      for (int64_t i = 0; i < axis_size; i++) {
+        axis_.push_back(axis[i]);
+        if (axis[i] < 0) {
+          axis_[i] += x_dim_size;
+        }
+      }
+      Tensor out_grad_ = out_grad;
+      if (out_grad.shape().size() != x.shape().size()) {
+        auto out_grad_shape = get_unsqueeze_dims(out_grad, axis_);
+        out_grad_ = reshape<T>(out_grad, out_grad_shape);
+      }
+      x_grad_tmp = expand<T>(out_grad_, IntArray(x_dim));
     }
-    x_grad_tmp = expand<T>(out_grad_, IntArray(x_dim));
-  }
 
-  set_output<T>(x_grad_tmp, x_grad);
+    set_output<T>(x_grad_tmp, x_grad);
+  }
 }
 
 template <typename T>
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -55,6 +55,7 @@
     "pd_op.elementwise_pow",
     "pd_op.softmax",
     "pd_op.matmul",
+    "pd_op.reduce_as",
 ]
 
 
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
@@ -123,6 +123,10 @@ def matmul_net(x, y):
     return paddle.matmul(x, y)
 
 
+def reduce_as_net(x, y):
+    return paddle.reduce_as(x, y)
+
+
 def apply_to_static(net, use_cinn, input_spec=None):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
@@ -1449,5 +1453,108 @@ def setUp(self):
         self.tol = 1e-5
 
 
+class TestPrimBaseOneGradTwoInputs(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.y_shape = [200, 40]
+        self.init_y_shape = [None, 200]
+        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.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = reduce_as_net
+        self.enable_cinn = False
+        self.tol = 1e-5
+        self.y_without_grad = True
+
+    def base_net(self, flag=None):
+        if flag == "prim":
+            core._set_prim_all_enabled(True)
+        x = paddle.to_tensor(self.x, stop_gradient=False)
+        y = paddle.to_tensor(self.y, stop_gradient=False)
+        if flag == "prim":
+            fn = apply_to_static(
+                self.net,
+                use_cinn=self.enable_cinn,
+                input_spec=[
+                    InputSpec(shape=self.init_x_shape, dtype='float32'),
+                    InputSpec(shape=self.init_y_shape, dtype='float32'),
+                ],
+            )
+            fn.train()
+        else:
+            fn = self.net
+        res = fn(x, y)
+        res.backward()
+        if self.y_without_grad:
+            grad = x.gradient()
+        else:
+            grad = y.gradient()
+        if flag == "prim":
+            core._set_prim_all_enabled(False)
+        return res, [grad]
+
+    def test_prim_all_dynamic(self):
+        res_ref, grad_ref = self.base_net()
+        res, grad = self.base_net("prim")
+
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+        for dr, d in zip(grad_ref, grad):
+            np.testing.assert_allclose(dr, d, rtol=self.tol, atol=self.tol)
+
+
+class TestPrimReduceAsWithGrad2(TestPrimBaseOneGradTwoInputs):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.y_shape = [30, 1, 40]
+        self.init_y_shape = [None, None, 40]
+        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.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = reduce_as_net
+        self.enable_cinn = False
+        self.tol = 1e-5
+        self.y_without_grad = True
+
+
+class TestPrimReduceAsWithGrad3(TestPrimBaseOneGradTwoInputs):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.y_shape = [30, 200, 1]
+        self.init_y_shape = [None, None, 1]
+        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.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = reduce_as_net
+        self.enable_cinn = False
+        self.tol = 1e-5
+        self.y_without_grad = True
+
+
+class TestPrimReduceAsWithGrad4(TestPrimBaseOneGradTwoInputs):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.y_shape = [30, 1, 1]
+        self.init_y_shape = [None, None, 1]
+        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.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = reduce_as_net
+        self.enable_cinn = False
+        self.tol = 1e-5
+        self.y_without_grad = True
+
+
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -55,6 +55,7 @@`
`55`	`55`	`"pd_op.elementwise_pow",`
`56`	`56`	`"pd_op.softmax",`
`57`	`57`	`"pd_op.matmul",`
	`58`	`+ "pd_op.reduce_as",`
`58`	`59`	`]`
`59`	`60`
`60`	`61`