(ctx, tensor)
| 27 | class AllGather(torch.autograd.Function): |
| 28 | @staticmethod |
| 29 | def forward(ctx, tensor): |
| 30 | ctx.rank = link.get_rank() |
| 31 | ctx.world_size = link.get_world_size() |
| 32 | |
| 33 | # y = tensor.new(ctx.world_size, *tensor.size()) |
| 34 | y = [tensor.new(*tensor.size()) for _ in range(ctx.world_size)] |
| 35 | link.allgather(y, tensor) |
| 36 | |
| 37 | y = torch.cat(y, 0).view(-1, *tensor.size()) |
| 38 | return y |
| 39 | |
| 40 | @staticmethod |
| 41 | def backward(ctx, grad_output): |
nothing calls this directly
no outgoing calls
no test coverage detected