| 610 | |
| 611 | |
| 612 | class _Scatter(Function): |
| 613 | def __init__(self, group=WORLD, device=None): |
| 614 | self.group = group |
| 615 | self.out_device = device |
| 616 | |
| 617 | def forward(self, data): |
| 618 | self.in_device = str(data.device) |
| 619 | return collective_comm( |
| 620 | data, CollectiveComm.Mode.SCATTER, self.group, self.out_device |
| 621 | ) |
| 622 | |
| 623 | def backward(self, grad): |
| 624 | # TODO backward with a part of grad |
| 625 | if grad is not None: |
| 626 | return gather(grad, self.group, self.in_device) |
| 627 | |
| 628 | |
| 629 | def scatter( |