def test_find_recvs():
z, recv_x, recv_x_plus_one, send_x, x_plus_one, from_node, send_x_plus_one = \
create_send_recv_graph()
assert set([recv_x]) == set(find_recvs(x_plus_one))
assert set([recv_x]) == set(find_recvs(recv_x))
assert len(find_recvs(from_node)) == 0
assert set([recv_x]) == set(find_recvs(send_x_plus_one))
assert set([recv_x_plus_one, recv_x]) == set(find_recvs(recv_x_plus_one))
assert set([recv_x_plus_one, recv_x]) == set(find_recvs(z))
def test_hetr_send_recv_graph_serialization():
"""
test serializing send/recv ops defined in comm_nodes for hetr communication
"""
z, recv_x, recv_x_plus_one, send_x, x_plus_one, from_node, send_x_plus_one = \
create_send_recv_graph()
ser_string = ser.serialize_graph([z])
py_graph = ser.deserialize_graph(ser_string)
orig_graph = Op.all_op_references([z])
for o1, o2 in zip(sorted(py_graph, key=lambda x: x.uuid),
sorted(orig_graph, key=lambda x: x.uuid)):
assert_object_equality(o1, o2)
def test_update_comm_deps():
with ng.metadata(transformer='cpu0'):
z, recv_x, recv_x_plus_one, send_x, x_plus_one, from_node, send_x_plus_one = \
create_send_recv_graph()
update_comm_deps((z, send_x))
assert recv_x_plus_one in z.all_deps