| 529 | |
| 530 | |
| 531 | class _Gather(Function): |
| 532 | def __init__(self, group=WORLD, device=None): |
| 533 | self.group = group |
| 534 | self.out_device = device |
| 535 | |
| 536 | def forward(self, data): |
| 537 | self.in_device = str(data.device) |
| 538 | return collective_comm( |
| 539 | data, CollectiveComm.Mode.GATHER, self.group, self.out_device |
| 540 | ) |
| 541 | |
| 542 | def backward(self, grad): |
| 543 | has_grad = _bcast_has_grad(self.group, grad) |
| 544 | if has_grad: |
| 545 | return scatter(grad, self.group, self.in_device) |
| 546 | |
| 547 | |
| 548 | def gather( |