| 791 | |
| 792 | |
| 793 | class _RemoteSend(Function): |
| 794 | def __init__(self, op: RemoteSend): |
| 795 | self.op = op |
| 796 | |
| 797 | def forward(self, data): |
| 798 | self.device = str(data.device) |
| 799 | (self.dummy,) = apply(self.op, data) |
| 800 | return self.dummy |
| 801 | |
| 802 | def backward(self, grad): |
| 803 | assert grad is None |
| 804 | has_grad = get_client().bcast_val(None, self.op.key, 2) |
| 805 | if has_grad: |
| 806 | return remote_recv(self.op.rank_to, device=self.device, inp=self.dummy,) |
| 807 | |
| 808 | |
| 809 | class _RemoteRecv(Function): |