fix assign kernel (#59609)

LiYuRio · web-flow · commit 787455ceaa85 · 2023-12-05T13:01:42.000+08:00
diff --git a/paddle/phi/core/distributed/auto_parallel/dist_attr.cc b/paddle/phi/core/distributed/auto_parallel/dist_attr.cc
@@ -268,6 +268,20 @@ bool TensorDistAttr::verify(const std::vector<int64_t>& tensor_shape) const {
   return true;
 }
 
+bool TensorDistAttr::verify_dynamic(
+    const std::vector<int64_t>& tensor_shape) const {
+  if (!verify_process_mesh(process_mesh_)) {
+    return false;
+  }
+  if (!verify_dims_mapping(dims_mapping_, tensor_shape)) {
+    return false;
+  }
+  if (!verify_partial_status()) {
+    return false;
+  }
+  return true;
+}
+
 std::string TensorDistAttr::to_string() const {
   std::string dist_str;
   dist_str += "{process_mesh: " + process_mesh_.to_string() + ", ";
diff --git a/paddle/phi/core/distributed/auto_parallel/dist_attr.h b/paddle/phi/core/distributed/auto_parallel/dist_attr.h
@@ -158,6 +158,7 @@ class TEST_API TensorDistAttr {
   bool verify_partial_status() const;
 
   bool verify(const std::vector<int64_t>& tensor_shape) const;
+  bool verify_dynamic(const std::vector<int64_t>& tensor_shape) const;
 
   // TensorDistAttr from_string(const std::string& dist_str);
   std::string to_string() const;
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.cc b/paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.cc
@@ -43,7 +43,7 @@ void ReshardFunction::SetValue(DistTensor* tensor, const DenseTensor& value) {
 void ReshardFunction::SetDistProps(DistTensor* tensor,
                                    const DDim& dims,
                                    const TensorDistAttr& dist_attr) {
-  PADDLE_ENFORCE_EQ(dist_attr.verify(common::vectorize(dims)),
+  PADDLE_ENFORCE_EQ(dist_attr.verify_dynamic(common::vectorize(dims)),
                     true,
                     phi::errors::InvalidArgument(
                         "The input dist_attr [%s] and dims [%s] are improper.",
@@ -56,10 +56,12 @@ void ReshardFunction::SetDistProps(DistTensor* tensor,
 
 void ReshardFunction::SetDistProps(DistTensor* tensor,
                                    const TensorDistAttr& dist_attr) {
-  PADDLE_ENFORCE_EQ(dist_attr.verify(common::vectorize(tensor->dims())),
+  PADDLE_ENFORCE_EQ(dist_attr.verify_dynamic(common::vectorize(tensor->dims())),
                     true,
                     phi::errors::InvalidArgument(
-                        "The input dist_attr and dims are improper."));
+                        "The input dist_attr [%s] and dims [%s] are improper.",
+                        dist_attr.to_string(),
+                        str_join(vectorize(tensor->dims()))));
 
   tensor->dist_attr_ = dist_attr;
 }
diff --git a/paddle/phi/kernels/assign_kernel.h b/paddle/phi/kernels/assign_kernel.h
@@ -34,7 +34,9 @@ DenseTensor Assign(const Context& dev_ctx, const DenseTensor& x) {
   MetaTensor meta_out(&out);
   MetaTensor meta_x(x);
   UnchangedInferMeta(meta_x, &meta_out);
-  AssignKernel<Context>(dev_ctx, x, &out);
+  if (x.initialized()) {
+    AssignKernel<Context>(dev_ctx, x, &out);
+  }
   return out;
 }
 
@@ -43,7 +45,9 @@ void Assign(const Context& dev_ctx, const DenseTensor& x, DenseTensor* out) {
   MetaTensor meta_out(out);
   MetaTensor meta_x(x);
   UnchangedInferMeta(meta_x, &meta_out);
-  AssignKernel<Context>(dev_ctx, x, out);
+  if (x.initialized()) {
+    AssignKernel<Context>(dev_ctx, x, out);
+  }
 }
 
 // In order to be compatible with the `AsDispensable` input in the original
