Diff Coverage

def _ensure_mesh_process_groups(mesh: DeviceMesh) -> None:
    """Lazily create per-axis process groups when mesh was built with ``init_backend=False``."""
    if hasattr(mesh, "_dim_group_names") and mesh._dim_group_names is not None:
        return
    mesh._dim_group_names = DeviceMesh._init_process_groups(  # pylint: disable=protected-access
        mesh._mesh_shape,
        mesh.mesh_dim_names,
        mesh._rank_list,
    )

    *,
    group_src: int = 0,
) -> Tensor:
    """Scatter tensor chunks along one mesh dimension (PyTorch ``mesh_scatter`` parity)."""
    _ensure_mesh_process_groups(mesh)
    group = mesh.get_group(mesh_dim)
    contiguous_list = [
        chunk.contiguous() if hasattr(chunk, "is_contiguous") and not chunk.is_contiguous() else chunk
        for chunk in scatter_list
    ]
    if platform.get_group_rank(group) == group_src:
        platform.scatter(output, list(contiguous_list), group=group, group_src=group_src)
    else:
        platform.scatter(output, None, group=group, group_src=group_src)
    return output


def mesh_broadcast(
    tensor: Tensor,

    *,
    group_src: int = 0,
) -> Tensor:
    """Broadcast a tensor along one mesh dimension (PyTorch ``mesh_broadcast`` parity)."""
    _ensure_mesh_process_groups(mesh)
    group = mesh.get_group(mesh_dim)
    if hasattr(tensor, "is_contiguous") and not tensor.is_contiguous():
        tensor = tensor.contiguous()
    platform.broadcast(tensor, group=group, group_src=group_src)
    return tensor

    local = tensor
    for mesh_dim, placement in enumerate(placements):
        if isinstance(placement, StridedShard):
            raise NotImplementedError(
                "distribute_tensor with src_data_rank does not support StridedShard yet; "
                "pass src_data_rank=None for local-only sharding."
            )
        if placement.is_shard():

        if placement.is_shard():
            shard_dim = _normalize_shard_dim(placement.dim, local.ndim)
            num_chunks = device_mesh.size(mesh_dim)
            if num_chunks <= 0:
                raise ValueError(f"invalid mesh dim size {num_chunks} on mesh_dim={mesh_dim}")
            chunks = tuple(local.chunk(num_chunks, dim=shard_dim))
            if not chunks:
                raise ValueError(f"cannot shard dim {shard_dim} into {num_chunks} chunks")
            output = platform.empty_like(chunks[0])
            local = mesh_scatter(output, chunks, device_mesh, mesh_dim, group_src=src_data_rank)
        elif placement.is_replicate() or placement.is_partial():
            local = mesh_broadcast(local, device_mesh, mesh_dim, group_src=src_data_rank)

            local = mesh_scatter(output, chunks, device_mesh, mesh_dim, group_src=src_data_rank)
        elif placement.is_replicate() or placement.is_partial():
            local = mesh_broadcast(local, device_mesh, mesh_dim, group_src=src_data_rank)
            if isinstance(placement, Partial):
                warnings.warn(
                    f"Partial placement {placement} during distribute_tensor: "
                    "broadcast only; partial partition is not applied yet.",
                    stacklevel=3,
                )

                    "broadcast only; partial partition is not applied yet.",
                    stacklevel=3,
                )
        else:
            raise RuntimeError(
                f"unsupported placement {placement} on device mesh dimension {mesh_dim}"
            )
    return local

        return data, handle

    @staticmethod
    def broadcast(data, src=None, group=None, async_op=False, group_src=None):
        if group_src is not None:
            ranks = MindSporePlatform.get_process_group_ranks(group)
            src = ranks[group_src]
        handle = dist.broadcast(data, src, group, async_op)
        if async_op:
            handle.wait()
        return data

        return data

    @staticmethod
    def scatter(output, scatter_list, src=None, group=None, async_op=False, group_src=None):
        group_name = group if isinstance(group, str) else getattr(group, "group_name", group)
        if group_src is not None:
            ranks = MindSporePlatform.get_process_group_ranks(group)
            src = ranks[group_src]
        if scatter_list is not None:
            scatter_list = [c.contiguous() if hasattr(c, "is_contiguous") and not c.is_contiguous() else c
                            for c in scatter_list]
        handle = dist.scatter(output, scatter_list, src, group_name, async_op=async_op)
        if async_op and handle is not None:
            handle.wait()
        return output

    @staticmethod
    def reduce_scatter_tensor(data, group_info, async_op=False):
        group_name = group_info if isinstance(group_info, str) else group_info.group_name

        return dist.get_group_rank(group, MindSporePlatform.get_rank())

    @staticmethod
    def get_group_rank(group=None) -> int:
        return MindSporePlatform.get_group_local_rank(group)

    @staticmethod
    def no_grad():
        return _no_grad()

    @staticmethod
    def get_group_rank(group):
        """Return this process's rank within *group*."""
        raise NotImplementedError("Platform subclasses must implement get_group_rank")

    @staticmethod
    def get_world_size():
        """Get the total number of processes in the default process group.

    @staticmethod
    def scatter(output, scatter_list, src=None, group=None, async_op=False, group_src=None):
        """Scatter tensor list from source rank to all ranks in group."""
        raise NotImplementedError("Platform subclasses must implement scatter")

    @staticmethod
    def isend(tensor, dst=None, group=None, tag=0):
        """Send tensor asynchronously to destination rank.

    @staticmethod
    def get_group_rank(group):
        """Return this process's rank within *group*."""
        return dist.get_group_rank(group, dist.get_rank())

    @staticmethod
    def get_world_size():
        """

        return data, handle

    @staticmethod
    def broadcast(data, src=None, group=None, async_op=False, group_src=None):
        if group_src is not None:
            handle = dist.broadcast(data, group=group, group_src=group_src, async_op=async_op)
        else:
            handle = dist.broadcast(data, src, group, async_op)
        if async_op and handle is not None:
            handle.wait()

    @staticmethod
    def scatter(output, scatter_list, src=None, group=None, async_op=False, group_src=None):
        if group_src is not None:
            handle = dist.scatter(
                output, scatter_list, group=group, group_src=group_src, async_op=async_op
            )
        else:
            handle = dist.scatter(output, scatter_list, src=src, group=group, async_op=async_op)
        if async_op and handle is not None:
            handle.wait()
        return output

    @staticmethod
    def isend(tensor, dst=None, group=None, tag=0):
        return dist.isend(tensor, dst, group, tag)