repeated_mesh_dims = mesh_dim
if reverse_repeated_shards:
# RNG follows the sequential placement application: the last
# placement subdivides each earlier shard and becomes major.
repeated_mesh_dims = list(reversed(repeated_mesh_dims))
rank_coord = [mesh_coordinate[d] for d in repeated_mesh_dims]
num_shards = [mesh_size[d] for d in repeated_mesh_dims]
# compute the shard idx and total number of shards
for idx, size in zip(rank_coord, num_shards):