def test_local_dimshuffle_alloc():
reshape_dimshuffle = out2in(local_dimshuffle_alloc)
x = tensor.vector("x")
out = tensor.alloc(x, 3, 2).dimshuffle("x", "x", 0, 1)
g = FunctionGraph([x], [out])
reshape_dimshuffle(g)
l = PerformLinker()
l.accept(g)
f = l.make_function()
assert f([3, 4]).ndim == 4
topo = g.toposort()
assert any([not isinstance(x, DimShuffle) for x in topo])
def test_0(self):
nodes = []
def wrap(fgraph, i, node, th):
nodes.append(node.op)
x, y, z = inputs()
e = mul(add(x, y), div(x, y))
fn, i, o = wrap_linker(FunctionGraph([x, y, z], [e]),
[PerformLinker(allow_gc=False)],
wrap).make_thunk()
i[0].data = 1
i[1].data = 2
fn()
assert nodes == [div, add, mul] or nodes == [add, div, mul]
assert o[0].data is None
def test_duallinker_mismatch():
x, y, z = inputs()
# bad_sub is correct in C but erroneous in Python
e = bad_sub(mul(x, y), mul(y, z))
g = Env([x, y, z], [e])
lnk = DualLinker(checker=_my_checker).accept(g)
fn = lnk.make_function()
# good
assert CLinker().accept(g).make_function()(1.0, 2.0, 3.0) == -4.0
# good
assert OpWiseCLinker().accept(g).make_function()(1.0, 2.0, 3.0) == -4.0
# (purposely) wrong
assert PerformLinker().accept(g).make_function()(1.0, 2.0, 3.0) == -10.0
with pytest.raises(MyExc):
# this runs OpWiseCLinker and PerformLinker in parallel and feeds
# variables of matching operations to _my_checker to verify that they
# are the same.
fn(1.0, 2.0, 3.0)
import theano
from theano import config, gof
from theano.compile.function.types import Supervisor
from theano.link.basic import PerformLinker
from theano.link.c.basic import CLinker, OpWiseCLinker
from theano.link.jax import JAXLinker
from theano.link.vm import VMLinker
_logger = logging.getLogger("theano.compile.mode")
# If a string is passed as the linker argument in the constructor for
# Mode, it will be used as the key to retrieve the real linker in this
# dictionary
predefined_linkers = {
"py": PerformLinker(), # Use allow_gc Theano flag
"c": CLinker(), # Don't support gc. so don't check allow_gc
"c|py": OpWiseCLinker(), # Use allow_gc Theano flag
"c|py_nogc": OpWiseCLinker(allow_gc=False),
"vm": VMLinker(use_cloop=False), # Use allow_gc Theano flag
"cvm": VMLinker(use_cloop=True), # Use allow_gc Theano flag
"vm_nogc": VMLinker(allow_gc=False, use_cloop=False),
"cvm_nogc": VMLinker(allow_gc=False, use_cloop=True),
"jax": JAXLinker(),
}
def register_linker(name, linker):
"""Add a `Linker` which can be referred to by `name` in `Mode`."""
if name in predefined_linkers:
raise ValueError(f"Linker name already taken: {name}")
def perform_linker(fgraph):
lnk = PerformLinker().accept(fgraph)
return lnk
def test_gc_never_pickles_temporaries():
x = tt.dvector()
r = x
for i in range(2): # TODO: 30 causes like LONG compilation due to MERGE
r = r + r / 10
optimizer = None
optimizer = "fast_run"
for f_linker, g_linker in [
(PerformLinker(allow_gc=True), PerformLinker(allow_gc=False)),
(OpWiseCLinker(allow_gc=True), OpWiseCLinker(allow_gc=False)),
]:
# f_linker has garbage collection
# g_linker has no garbage collection
f = theano.function([x],
r,
mode=theano.Mode(optimizer=optimizer,
linker=f_linker))
g = theano.function([x],
r,
mode=theano.Mode(optimizer=optimizer,
linker=g_linker))
pre_f = pickle.dumps(f)
# pre_g = pickle.dumps(g)
len_pre_f = len(pre_f)
# len_pre_g = len(pre_g)
# We can't compare the content or the length of the string
# between f and g. 2 reason, we store some timming information
# in float. They won't be the same each time. Different float
# can have different length when printed.
def a(fn):
return len(pickle.dumps(fn.maker))
assert a(f) == a(f) # some sanity checks on the pickling mechanism
assert a(g) == a(g) # some sanity checks on the pickling mechanism
def b(fn):
return len(
pickle.dumps(
theano.compile.function.types._pickle_Function(fn)))
assert b(f) == b(f) # some sanity checks on the pickling mechanism
def c(fn):
return len(pickle.dumps(fn))
assert c(f) == c(f) # some sanity checks on the pickling mechanism
assert c(g) == c(g) # some sanity checks on the pickling mechanism
# now run the function once to create temporaries within the no-gc
# linker
f(np.ones(100, dtype="float64"))
g(np.ones(100, dtype="float64"))
# serialize the functions again
post_f = pickle.dumps(f)
post_g = pickle.dumps(g)
len_post_f = len(post_f)
len_post_g = len(post_g)
# assert that f() didn't cause the function to grow
# allow_gc should leave the function un-changed by calling
assert len_pre_f == len_post_f, (len_pre_f, len_post_f)
# assert that g() didn't cause g to grow because temporaries
# that weren't collected shouldn't be pickled anyway
# Allow for a couple of bytes of difference, since timing info,
# for instance, can be represented as text of varying size.
assert abs(len_post_f - len_post_g) < 256, (f_linker, len_post_f,
len_post_g)