Function broadcast

imperative/python/megengine/distributed/functional.py:201–244

r"""Broadcast tensor data from root process to others. Args: inp: Input tensor. group: The process group to work on. The default group is WORLD which means all processes available. You can use a list of process ranks to create new group to work on it, e.g

(
    inp: Tensor, group: Optional[Group] = WORLD, device: Optional[str] = None,
)

def broadcast(
    inp: Tensor, group: Optional[Group] = WORLD, device: Optional[str] = None,
) -> Tensor:
    r"""Broadcast tensor data from root process to others.

    Args:
        inp: Input tensor.
        group: The process group to work on.
            The default group is WORLD which means all processes available.
            You can use a list of process ranks to create new group to work on it, e.g. [1, 3, 5].
        device: The specific device to execute this operator.
            None default device means the device of inp will be used.
            Specify "gpu0:1" to execute this operator on diffrent cuda stream,
            1 is stream id, and default stream id is 0.

    Returns:
        Result tensor.

    Examples:

        .. code-block::

           input = Tensor([rank])
           # Rank 0 # input: Tensor([0])
           # Rank 1 # input: Tensor([1])
           output = broadcast(input)
           # Rank 0 # output: Tensor([0])
           # Rank 1 # output: Tensor([0])

           input = Tensor([rank])
           group = Group([1, 0]) # first rank is root
           output = broadcast(input, group)
           # Rank 0 # output: Tensor([1])
           # Rank 1 # output: Tensor([1])
    """
    shape, dtype = _bcast_shape_dtype(group, inp)
    if group.rank != 0:
        # dummy input to infer shape
        inp = _dummy_input(shape, dtype, device)

    _bcast_tracer_state(group, inp)

    out = _Broadcast(group, device)(inp)
    return out


def _bcast_param(