r"""Split tensor in root process at first dimension. 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, axis=0,
)
| 627 | |
| 628 | |
| 629 | def scatter( |
| 630 | inp: Tensor, group: Optional[Group] = WORLD, device: Optional[str] = None, axis=0, |
| 631 | ) -> Tensor: |
| 632 | r"""Split tensor in root process at first dimension. |
| 633 | |
| 634 | Args: |
| 635 | inp: Input tensor. |
| 636 | group: The process group to work on. |
| 637 | The default group is WORLD which means all processes available. |
| 638 | You can use a list of process ranks to create new group to work on it, e.g. [1, 3, 5]. |
| 639 | device: The specific device to execute this operator. |
| 640 | None default device means the device of inp will be used. |
| 641 | Specify "gpu0:1" to execute this operator on diffrent cuda stream, |
| 642 | 1 is stream id, and default stream id is 0. |
| 643 | axis: The concat axis for collective_comm result |
| 644 | The default axis is 0 |
| 645 | |
| 646 | Returns: |
| 647 | Split tensor. |
| 648 | |
| 649 | Examples: |
| 650 | |
| 651 | .. code-block:: |
| 652 | |
| 653 | input = Tensor([0 1]) + rank*2 |
| 654 | # Rank 0 # input: Tensor([0 1]) |
| 655 | # Rank 1 # input: Tensor([2 3]) |
| 656 | output = scatter(input) |
| 657 | # Rank 0 # output: Tensor([0]) |
| 658 | # Rank 1 # output: Tensor([1]) |
| 659 | |
| 660 | input = Tensor([0 1]) + rank*2 |
| 661 | group = Group([1, 0]) # first rank is root |
| 662 | output = scatter(input, group) |
| 663 | # Rank 0 # output: Tensor([3]) |
| 664 | # Rank 1 # output: Tensor([2]) |
| 665 | """ |
| 666 | shape, dtype = _bcast_shape_dtype(group, inp) |
| 667 | if group.rank != 0: |
| 668 | # dummy input to infer shape |
| 669 | inp = _dummy_input(shape, dtype, device) |
| 670 | |
| 671 | _bcast_tracer_state(group, inp) |
| 672 | |
| 673 | assert ( |
| 674 | list(inp._tuple_shape)[axis] % group.size == 0 |
| 675 | ), "current axis: {} can't devided by group size".format(axis) |
| 676 | |
| 677 | if axis != 0: |
| 678 | group_size = group.size |
| 679 | k_new_shape = list(inp._tuple_shape) |
| 680 | k_new_shape[axis] //= group_size |
| 681 | k_new_shape[0] *= group_size |
| 682 | new_shape = list(inp._tuple_shape) |
| 683 | new_shape[axis] //= group_size |
| 684 | new_shape.insert(axis, group_size) |
| 685 | index = ( |
| 686 | [axis] |