def make_scheduler(*sched_cmps):
from theano.gof.sched import sort_schedule_fn
from theano.tensor.io import mpi_cmps
if theano.config.device == 'gpu':
from theano.sandbox.cuda. async import gpu_cmps
else:
gpu_cmps = ()
return sort_schedule_fn(*(tuple(sched_cmps) + mpi_cmps + gpu_cmps))
def make_scheduler(*sched_cmps):
from theano.gof.sched import sort_schedule_fn
from theano.tensor.io import mpi_cmps
if theano.config.device == "gpu":
from theano.sandbox.cuda.async import gpu_cmps
else:
gpu_cmps = ()
return sort_schedule_fn(*(tuple(sched_cmps) + mpi_cmps + gpu_cmps))
def test_sort_schedule_fn():
import theano
from theano.gof.sched import sort_schedule_fn, make_depends
x = theano.tensor.matrix('x')
y = theano.tensor.dot(x[:5]*2, x.T+1).T
str_cmp = lambda a, b: cmp(str(a), str(b)) # lexicographical sort
linker = theano.OpWiseCLinker(schedule=sort_schedule_fn(str_cmp))
mode = theano.Mode(linker=linker)
f = theano.function((x,), (y,), mode=mode)
nodes = f.maker.linker.make_all()[-1]
depends = make_depends()
for a, b in zip(nodes[:-1], nodes[1:]):
if not depends((b, a)):
assert str(a) < str(b)
def test_sort_schedule_fn():
import theano
from theano.gof.sched import sort_schedule_fn, make_depends
x = theano.tensor.matrix('x')
y = theano.tensor.dot(x[:5] * 2, x.T + 1).T
str_cmp = lambda a, b: cmp(str(a), str(b)) # lexicographical sort
linker = theano.OpWiseCLinker(schedule=sort_schedule_fn(str_cmp))
mode = theano.Mode(linker=linker)
f = theano.function((x, ), (y, ), mode=mode)
nodes = f.maker.linker.make_all()[-1]
depends = make_depends()
for a, b in zip(nodes[:-1], nodes[1:]):
if not depends((b, a)):
assert str(a) < str(b)
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
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 = theano.Mode(optimizer=None,
linker=theano.OpWiseCLinker(schedule=scheduler))
with change_flags(compute_test_value="off"):
if rank == 0:
x = theano.tensor.matrix("x", dtype=dtype)
y = x + 1
send_request = send(y, 1, 11)
z = recv(shape, dtype, 1, 12)
f = theano.function([x], [send_request, z], mode=mode)
xx = np.random.rand(*shape).astype(dtype)
expected = (xx + 1) * 2
from __future__ import absolute_import, print_function, division
import os
import subprocess
from nose.plugins.skip import SkipTest
import theano
from theano import change_flags
from theano.compat import PY3
from theano.gof.sched import sort_schedule_fn
from theano.tensor.io import (send, recv, mpi_cmps, MPISend, MPISendWait,
mpi_send_wait_cmp, mpi_tag_cmp, mpi_enabled)
mpi_scheduler = sort_schedule_fn(*mpi_cmps)
mpi_linker = theano.OpWiseCLinker(schedule=mpi_scheduler)
mpi_mode = theano.Mode(linker=mpi_linker)
@change_flags(compute_test_value='off')
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_send():
x = theano.tensor.matrix('x')
