def test_mpi_tag_ordering():
x = recv((2, 2), "float32", 1, 12)
y = recv((2, 2), "float32", 1, 11)
z = recv((2, 2), "float32", 1, 13)
f = aesara.function([], [x, y, z], mode=mpi_mode)
nodes = f.maker.linker.make_all()[-1]
assert all(node.op.tag == tag for node, tag in zip(nodes, (11, 12, 13, 11, 12, 13)))
def test_recv():
x = recv((10, 10), "float64", 0, 11)
assert x.dtype == "float64"
assert x.broadcastable == (False, False)
recvnode = x.owner.inputs[0].owner
assert recvnode.op.source == 0
assert recvnode.op.tag == 11
def test_can_make_function():
x = recv((5, 5), "float32", 0, 11)
y = x + 1
assert aesara.function([], [y])
shape = (2, 2)
dtype = "float32"
scheduler = sort_schedule_fn(*mpi_cmps)
mode = aesara.compile.mode.Mode(
optimizer=None,
linker=aesara.link.c.basic.OpWiseCLinker(schedule=scheduler))
with config.change_flags(compute_test_value="off"):
if rank == 0:
x = matrix("x", dtype=dtype)
y = x + 1
send_request = send(y, 1, 11)
z = recv(shape, dtype, 1, 12)
f = aesara.function([x], [send_request, z], mode=mode)
xx = np.random.random(shape).astype(dtype)
expected = (xx + 1) * 2
_, zz = f(xx)
same = np.linalg.norm(zz - expected) < 0.001
# The parent test will look for "True" in the output
stdout.write(str(same))
if rank == 1:
y = recv(shape, dtype, 0, 11)