def set_aesara_flags(): with config.change_flags(cxx="", compute_test_value="raise"): yield
if size != 2: stderr.write("mpiexec failed to create a world with two nodes.\n" "Closing with success message.") stdout.write("True") exit(0) 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.rand(*shape).astype(dtype) expected = (xx + 1) * 2 _, zz = f(xx) same = np.linalg.norm(zz - expected) < 0.001