MPI

dragon.operators.mpi.MPIBroadcast(inputs, root, mpi_ranks=None, **kwargs)

Broadcast a tensor to all nodes in the MPIGroup.

Parameters:
  • inputs (Tensor) – The tensor to broadcast.
  • root (int) – The world rank of root node.
  • mpi_ranks (sequence of int, optional) – The world rank of nodes in group. Default is None (Use All).
Returns:

The broadcast output.

Return type:

Tensor

Notes

For root, the output shares the input.

For others, the input is inaccessible.

dragon.operators.mpi.MPIGather(inputs, root, mpi_ranks=None, **kwargs)

Gather a tensor from all nodes to root in the MPIGroup.

Parameters:
  • inputs (Tensor) – The tensor to gather.
  • root (int) – The world rank of root node.
  • mpi_ranks (sequence of int, optional) – The world rank of nodes in group. Default is None (Use All).
Returns:

The gathered outputs.

Return type:

sequence of Tensor

Notes

The number of outputs is decided on the number of mpi_ranks.

The outputs are accessible only for root and vice versa.