import paddle
import paddle.distributed as dist
a = paddle.to_tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['x', 'y'])
placements = [
    dist.Shard(0, co_shard_order=0),
    dist.Shard(0, co_shard_order=1),
]
b = dist.shard_tensor(a, mesh, placements)
print(b.placements)  # [Shard(0, shard_order=0), Shard(0, shard_order=1)]
print(b._local_value()) # rank0 [[1, 2]], rank1 [[3, 4]], rank2 [[5, 6]], rank3 [[7, 8]]

import paddle
import paddle.distributed as dist
a = paddle.to_tensor([[1, 2], [3, 4], [5, 6], [7, 8]], dtype='float32')
mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['x', 'y'])
placements = [dist.Shard(0), dist.Shard(1)]
input = dist.shard_tensor(a, mesh, placements)
out = paddle.reshape(input, [-1])
print(out.placements) # [Shard(0, shard_order=0), Shard(0, shard_order=1)]

import paddle
import paddle.distributed as dist
mesh = dist.ProcessMesh([[0, 1], [2, 3]], dim_names=['x', 'y'])
a = paddle.to_tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
placements = [dist.Shard(0, split_factor=2), dist.Replicate()]
b = dist.shard_tensor(a, mesh, placements)
print(b.placements) # [Shard(0, split_factor=2), Replicate()]
print(b._local_value()) # rank0 rank1 [[1, 2], [5, 6]] , ran2 rank3 [[3, 4], [7, 8]]