import paddle
import paddle.distributed as dist
process_mesh = dist.ProcessMesh([0, 1], dim_names=['mp'])
# before this PR, only support
linear1 = paddle.nn.Linear(100, 200)
linear1.weight = dist.shard_tensor(linear1.weight, process_mesh, [dist.Shard(1)]) # can shard input data and model parameter
noise = paddle.randn(x.shape)
noise = dist.shard_tensor(noise, process_mesh, [dist.Replicate()) # In other cases, placements must be all Replicate, otherwise an error will be reported
# after this PR, also support
timesteps = paddle.randint(0, self.num_timesteps, (x.shape[0],))
timesteps = dist.shard_tensor(timesteps, process_mesh, [dist.Shard(1))  # can shard randomly sampled data(int)
noise = paddle.randn(x.shape)
noise = dist.shard_tensor(noise, process_mesh, [dist.Shard(1))  # can shard randomly sampled data(float)