fix concat spmd rule when input is one dimensional

[jeff41404]

PR Category

Auto Parallel

PR Types

Bug fixes

Description

pcard-89497
when use distributed auto parallel, executing concat error on 1D tensors can be reproduced by the following code. this PR will fix this issue. By the way, fix issue of concat error when put_along_axis use concat in dynamic shape.

# python -m paddle.distributed.launch --gpus=0,1,2,3,4,5,6,7 --log_dir logs xxx.py
import paddle
import paddle.distributed as dist

paddle.seed(42)
x = paddle.randn(shape=[16], dtype='float32')
y = paddle.randn(shape=[8], dtype='float32')
z = paddle.concat([x, y]) # ok, shape is [24]

mesh = dist.ProcessMesh([0, 1, 2, 3, 4, 5, 6, 7], dim_names=['dp'])
dist_x = dist.shard_tensor(x, mesh, [dist.Replicate()])
# dist_x = dist.shard_tensor(x, mesh, [dist.Shard(0)]) # shard x/y will get the same result
dist_y = dist.shard_tensor(y, mesh, [dist.Replicate()])
dist_z = paddle.concat([dist_x, dist_y]) # will report error

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[From00]

LGTM

[zhiqiu]

LGTM