moose
moose copied to clipboard
GrpcRuntime: only send arguments to annotated placements
In the eDSL we have the placement where an argument is coming from:
def simple_computation(
x: pm.Argument(placement=alice, vtype=pm.TensorType(pm.float64)),
y: pm.Argument(placement=bob, vtype=pm.TensorType(pm.float64)),
):
which means we might be able to only send x
to Alice, and only y
to Bob, instead of currently sending everything to everyone.