def create_send_recv_graph():
ax_a = ng.make_axis(length=10, name='A')
ax_b = ng.make_axis(length=15, name='B')
axes = ng.make_axes([ax_a, ax_b])
with ng.metadata(device=None,
device_id=None,
transformer=None,
host_transformer=None):
from_node = ng.placeholder(axes)
to_node = ng.placeholder(axes)
send_x = SendOp(from_node=from_node)
recv_x = RecvOp(to_node=to_node, send_node=send_x)
with ng.metadata(parallel=ax_a,
device=None,
device_id=None,
transformer=None,
host_transformer=None):
x_plus_one = recv_x + 1
send_x_plus_one = SendOp(from_node=x_plus_one)
recv_x_plus_one = RecvOp(to_node=to_node, send_node=send_x_plus_one)
with ng.metadata(device=None,
device_id=None,
transformer=None,
host_transformer=None):
z = recv_x_plus_one + 2
return z, recv_x, recv_x_plus_one, send_x, x_plus_one, from_node, send_x_plus_one
def test_comm_path_exists():
axes = ng.make_axes([ng.make_axis(length=10, name='A'), ng.make_axis(length=15, name='B')])
with ng.metadata(device=None, device_id=None, transformer=None, host_transformer=None):
from_node = ng.placeholder(axes)
to_node = ng.placeholder(axes)
send_x = SendOp(from_node=from_node)
recv_x = RecvOp(to_node=to_node, send_node=send_x)
with ng.metadata(device=None, device_id=None, transformer=None, host_transformer=None):
x_plus_one = recv_x + 1
assert comm_path_exists(recv_x, send_x)
assert comm_path_exists(x_plus_one, send_x)