[Auto Parallel] add more tests for expand spmd (#73154)

ooooo-create · web-flow · commit 700f6e83aac6 · 2025-06-13T10:16:30.000+08:00
* add a test in ci

* fix bugs

* use set to partial_dims check
diff --git a/paddle/phi/infermeta/spmd_rules/expand.cc b/paddle/phi/infermeta/spmd_rules/expand.cc
@@ -41,23 +41,27 @@ SpmdInfo ExpandGradInferSpmd(const DistMetaTensor& x,
                              const IntArray& shape) {
   EXTRACT_SHAPE_AND_DIST_ATTR(x);
   EXTRACT_SHAPE_AND_DIST_ATTR(out_grad);
-  if (x_shape.size() == out_grad_shape.size()) {
-    return {{x_dist_attr_src, out_grad_dist_attr_src}, {x_dist_attr_src}};
-  }
   size_t axis =
       std::abs(static_cast<int>(out_grad.dims().size() - x.dims().size()));
-  std::vector<int64_t> x_grad_dims_mapping;
+  std::vector<int64_t> x_grad_dims_mapping(x_ndim, -1);
+  std::vector<int64_t> partial_dims;
   for (size_t i = 0; i < out_grad_dims_mapping_src.size(); ++i) {
     if (i < axis || i >= axis + x.dims().size() ||
         out_grad.dims()[i] != x.dims()[i - axis]) {
+      if (out_grad_dims_mapping_src[i] >= 0) {
+        partial_dims.push_back(out_grad_dims_mapping_src[i]);
+      }
       continue;
     }
-    x_grad_dims_mapping.push_back(out_grad_dims_mapping_src[i]);
+    x_grad_dims_mapping[i - axis] = out_grad_dims_mapping_src[i];
   }
   TensorDistAttr x_grad_dist_attr =
       CopyTensorDistAttrForOutput(x_dist_attr_src);
+  TensorDistAttr x_dist_attr = x_grad_dist_attr;
+  x_dist_attr.set_dims_mapping(x_grad_dims_mapping);
   x_grad_dist_attr.set_dims_mapping(x_grad_dims_mapping);
-  return {{x_dist_attr_src, out_grad_dist_attr_src}, {x_grad_dist_attr}};
+  x_grad_dist_attr.set_partial_status(partial_dims);
+  return {{x_dist_attr, out_grad_dist_attr_src}, {x_grad_dist_attr}};
 }
 
 }  // namespace phi::distributed
diff --git a/test/cpp/auto_parallel/expand_spmd_rule_test.cc b/test/cpp/auto_parallel/expand_spmd_rule_test.cc
@@ -73,6 +73,31 @@ TEST(ExpandInferSpmd, Ctor) {
   EXPECT_EQ(get_dims_mapping(spmdinfo3.first[1]),
             std::vector<int64_t>({-1, 1}));
   EXPECT_EQ(get_dims_mapping(spmdinfo3.second[0]), std::vector<int64_t>({1}));
+
+  // Test case backward 2: ExpandGrad with shape {2, 2, -1}
+  auto x4 = CreateDistMetaTensor({1, 8}, {-1, 1}, process_mesh);
+  auto out4 = CreateDistMetaTensor({2, 2, 8}, {-1, -1, 1}, process_mesh);
+  phi::IntArray shape4 = {2, 2, -1};
+  auto spmdinfo4 = ExpandGradInferSpmd(x4, out4, shape4);
+  EXPECT_EQ(get_dims_mapping(spmdinfo4.first[0]),
+            std::vector<int64_t>({-1, 1}));
+  EXPECT_EQ(get_dims_mapping(spmdinfo4.first[1]),
+            std::vector<int64_t>({-1, -1, 1}));
+  EXPECT_EQ(get_dims_mapping(spmdinfo4.second[0]),
+            std::vector<int64_t>({-1, 1}));
+
+  // Test case backward 3: ExpandGrad with shape {2, 2, -1}
+  auto x5 = CreateDistMetaTensor({1, 8}, {-1, 1}, process_mesh);
+  auto out5 = CreateDistMetaTensor({2, 2, 8}, {-1, 0, 1}, process_mesh);
+  phi::IntArray shape5 = {2, 2, -1};
+  auto spmdinfo5 = ExpandGradInferSpmd(x5, out5, shape5);
+  EXPECT_EQ(get_dims_mapping(spmdinfo5.first[0]),
+            std::vector<int64_t>({-1, 1}));
+  EXPECT_EQ(get_dims_mapping(spmdinfo5.first[1]),
+            std::vector<int64_t>({-1, 0, 1}));
+  EXPECT_EQ(get_dims_mapping(spmdinfo5.second[0]),
+            std::vector<int64_t>({-1, 1}));
+  EXPECT_EQ(get_partial_dims(spmdinfo5.second[0]), std::set<int64_t>({0}));
 }
 
 }  // namespace auto_parallel
