r"""Gather tensors across the specified group. Only root process will receive the final result. 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
(
inp: Tensor, group: Optional[Group] = WORLD, device: Optional[str] = None, axis=0,
)
| 546 | |
| 547 | |
| 548 | def gather( |
| 549 | inp: Tensor, group: Optional[Group] = WORLD, device: Optional[str] = None, axis=0, |
| 550 | ) -> Tensor: |
| 551 | r"""Gather tensors across the specified group. |
| 552 | Only root process will receive the final result. |
| 553 | |
| 554 | Args: |
| 555 | inp: Input tensor. |
| 556 | group: The process group to work on. |
| 557 | The default group is WORLD which means all processes available. |
| 558 | You can use a list of process ranks to create new group to work on it, e.g. [1, 3, 5]. |
| 559 | device: The specific device to execute this operator. |
| 560 | None default device means the device of inp will be used. |
| 561 | Specify "gpu0:1" to execute this operator on diffrent cuda stream, |
| 562 | 1 is stream id, and default stream id is 0. |
| 563 | axis: The concat axis for collective_comm result |
| 564 | |
| 565 | Examples: |
| 566 | |
| 567 | .. code-block:: |
| 568 | |
| 569 | input = Tensor([rank]) |
| 570 | # Rank 0 # input: Tensor([0]) |
| 571 | # Rank 1 # input: Tensor([1]) |
| 572 | output = gather(input) |
| 573 | # Rank 0 # output: Tensor([0 1]) |
| 574 | # Rank 1 # output: None |
| 575 | |
| 576 | input = Tensor([rank]) |
| 577 | group = Group([1, 0]) # first rank is root |
| 578 | output = gather(input, group) |
| 579 | # Rank 0 # output: None |
| 580 | # Rank 1 # output: Tensor([1 0]) |
| 581 | """ |
| 582 | assert ( |
| 583 | axis < inp.ndim |
| 584 | ), "your concat_axis exceeds the dim of the tensor, the tensor shape is {}".format( |
| 585 | inp.shape |
| 586 | ) |
| 587 | |
| 588 | out = _Gather(group, device)(inp) |
| 589 | |
| 590 | if group.rank == 0: |
| 591 | if axis == 0: |
| 592 | return out |
| 593 | else: |
| 594 | group_size = group.size |
| 595 | transformed_shape = list(inp._tuple_shape) |
| 596 | transformed_shape[axis] *= group_size |
| 597 | n, *shp = out._tuple_shape |
| 598 | index = ( |
| 599 | [_ for _ in range(1, axis)] |
| 600 | + [axis, 0] |
| 601 | + [_ for _ in range(axis + 1, out.ndim + 1)] |
| 602 | ) |
| 603 | return ( |
| 604 | out.reshape(group_size, n // group_size, *shp) |
| 605 | .transpose(index) |