def recv(self, reduce_func, msg, recv_mode="dst"):
if recv_mode != "dst":
raise ValueError(
"Currently DistGPUGraph can only support recv_mode=='dst'")
output = super(DistGPUGraph, self).recv(
msg=msg, reduce_func=reduce_func, recv_mode=recv_mode)
out = op.all_reduce_sum_with_grad(output)
return out
def send_recv(self, feature, reduce_func="sum"):
output = super(DistGPUGraph, self).send_recv(feature=feature,
reduce_func="sum")
output = op.all_reduce_sum_with_grad(output)
return output
def outdegree(self, nodes=None):
degree = super(DistGPUGraph, self).outdegree(nodes=nodes)
degree = op.all_reduce_sum_with_grad(degree)
return degree
