fix concat spmd rule when input is one dimensional (#71253)

jeff41404 · Glencsa · web-flow · commit fa16d0da4ef6 · 2025-02-26T23:53:48.000+08:00
* fix concat spmd rule when input is one dimensional

* add unit test case

* fix issue of coverage and add unit test case

---------

Co-authored-by: Glencsa &lt;3501406249@qq.com&gt;
diff --git a/paddle/phi/infermeta/spmd_rules/concat.cc b/paddle/phi/infermeta/spmd_rules/concat.cc
@@ -85,8 +85,17 @@ SpmdInfo ConcatInferSpmd(const std::vector<DistMetaTensor>& x, int axis) {
   std::string align_axis;
   std::tie(all_axis, align_axis) = FillConcatNotation(ndim, dim);
   std::vector<std::string> axis_names(input_attrs.size(), all_axis);
-  AlignDimsSharding(
-      &input_attrs, tensor_shapes, axis_names, {}, align_axis, true);
+  if (ndim == 1 && align_axis.empty()) {
+    // Simply set the 1D tensor to Replicate, and calling AlignDimsSharding
+    // requires !align_axis.empty()
+    std::vector<int64_t> dims_mapping(1, -1);
+    for (size_t i = 0; i < input_attrs.size(); i++) {
+      input_attrs[i].set_dims_mapping(dims_mapping);
+    }
+  } else {
+    AlignDimsSharding(
+        &input_attrs, tensor_shapes, axis_names, {}, align_axis, true);
+  }
 
   auto out_dist_attr =
       CopyTensorDistAttrForOutput(input_attrs[non_empty_index]);
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -6462,11 +6462,26 @@ def infer_dynamic_broadcast_shape(
     Returns:
         Tensor: The shape tensor for later broadcasting
     """
-    new_shapes = [
-        arr_shape[:axis],
-        indices_shape[axis : axis + 1],
-        arr_shape[axis + 1 :],
-    ]
+    if axis == 0:
+        new_shapes = [
+            indices_shape[
+                :1
+            ],  # use indices_shape[0] will error in concat after, because its shape is [], and shape of arr_shape[1:] is [1]
+            arr_shape[1:],
+        ]
+    elif axis == arr_shape_dim - 1:
+        new_shapes = [
+            arr_shape[:axis],
+            indices_shape[
+                axis:
+            ],  # use indices_shape[axis] will error in concat after, because its shape is [], and shape of arr_shape[axis:] is [1]
+        ]
+    else:
+        new_shapes = [
+            arr_shape[:axis],
+            indices_shape[axis : axis + 1],
+            arr_shape[axis + 1 :],
+        ]
     return paddle.concat(new_shapes)
 
 
diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc
@@ -752,6 +752,29 @@ TEST(ConcatRule, Ctor) {
   }
   check_dim_mapping(inferred_dist_attrs.second[0], {1, -1, 0});
   check_partial_dims(inferred_dist_attrs.second[0], {});
+
+  // test 3，special case: concat one dimensional tensor
+  shapes = {{16}, {32}, {64}};
+  dim_mappings = {{0}, {1}, {-1}};
+  partial_status = {{}, {}, {1}};
+  inputs = build_inputs();
+  inferred_dist_attrs = phi::distributed::ConcatInferSpmd(inputs, 0);
+  // list of tensor => single tensor
+  EXPECT_EQ(inferred_dist_attrs.first.size(), static_cast<size_t>(1));
+  EXPECT_EQ(inferred_dist_attrs.second.size(), static_cast<size_t>(1));
+  EXPECT_TRUE(
+      paddle::holds_alternative<std::vector<phi::distributed::TensorDistAttr>>(
+          inferred_dist_attrs.first[0]));
+  EXPECT_TRUE(paddle::holds_alternative<phi::distributed::TensorDistAttr>(
+      inferred_dist_attrs.second[0]));
+  auto& inputs_infer3 = PADDLE_GET_CONST(std::vector<TensorDistAttr>,
+                                         inferred_dist_attrs.first[0]);
+  for (auto e : inputs_infer3) {
+    check_dim_mapping(e, {-1});
+    check_partial_dims(e, {});
+  }
+  check_dim_mapping(inferred_dist_attrs.second[0], {-1});
+  check_partial_dims(inferred_dist_attrs.second[0], {});
 }
 
 TEST(StackRule, Ctor) {
diff --git a/test/legacy_test/test_put_along_axis_op.py b/test/legacy_test/test_put_along_axis_op.py
@@ -27,10 +27,10 @@
 paddle.enable_static()
 
 
-def put_along_axis_net(arr):
+def put_along_axis_net(arr, axis=-1):
     indices = paddle.to_tensor([[[[2]]]], dtype='int32', stop_gradient=False)
-    return paddle.tensor.put_along_axis(
-        arr, indices=indices, values=-4.0, axis=-2, reduce='add'
+    return paddle.put_along_axis(
+        arr, indices=indices, values=-4.0, axis=axis, reduce='add'
     )
 
 
@@ -1346,6 +1346,7 @@ def setUp(self):
         self.enable_cinn = False
         self.tol = 1e-6
         self.dtype = "float32"
+        self.axis = -2
         self.input_specs = [
             InputSpec(
                 shape=(-1, -1, -1, -1),
@@ -1370,7 +1371,7 @@ def train(self, to_static):
         else:
             net = self.net
 
-        res = net(arr)
+        res = net(arr, self.axis)
         res.backward()
         arr_grad = arr.gradient()
         return res, arr_grad
@@ -1389,6 +1390,60 @@ def test_dynamic_static(self):
                 np.testing.assert_allclose(dr, d, rtol=self.tol, atol=self.tol)
 
 
+class TestPutAlongAxisDynamicShape1(TestPutAlongAxisDynamicShape):
+    def setUp(self):
+        np.random.seed(2024)
+        self.net = put_along_axis_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+        self.dtype = "float32"
+        self.axis = 0
+        self.input_specs = [
+            InputSpec(
+                shape=(-1, -1, -1, -1),
+                dtype=self.dtype,
+                stop_gradient=False,
+            )
+        ]
+        self.arr = np.random.random([16, 16, 16, 16]).astype(self.dtype)
+
+
+class TestPutAlongAxisDynamicShape2(TestPutAlongAxisDynamicShape):
+    def setUp(self):
+        np.random.seed(2024)
+        self.net = put_along_axis_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+        self.dtype = "float32"
+        self.axis = -1
+        self.input_specs = [
+            InputSpec(
+                shape=(-1, -1, -1, -1),
+                dtype=self.dtype,
+                stop_gradient=False,
+            )
+        ]
+        self.arr = np.random.random([20, 20, 20, 20]).astype(self.dtype)
+
+
+class TestPutAlongAxisDynamicShape3(TestPutAlongAxisDynamicShape):
+    def setUp(self):
+        np.random.seed(2024)
+        self.net = put_along_axis_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+        self.dtype = "float32"
+        self.axis = 3
+        self.input_specs = [
+            InputSpec(
+                shape=(-1, -1, -1, -1),
+                dtype=self.dtype,
+                stop_gradient=False,
+            )
+        ]
+        self.arr = np.random.random([32, 32, 32, 32]).astype(self.dtype)
+
+
 if __name__ == "__main__":
     paddle.enable_static()
     unittest.main()
