def test_optimize_graph():
x, y = vectors("xy")
@optimizer
def custom_opt(fgraph):
fgraph.replace(x, y, import_missing=True)
x_opt = optimize_graph(x, custom_opt=custom_opt)
assert x_opt is y
x_opt = optimize_graph(FunctionGraph(outputs=[x], clone=False),
custom_opt=custom_opt)
assert x_opt.outputs[0] is y
def test_KanrenRelationSub_dot():
"""Make sure we can run miniKanren "optimizations" over a graph until a fixed-point/normal-form is reached."""
x_at = at.vector("x")
c_at = at.vector("c")
d_at = at.vector("d")
A_at = at.matrix("A")
B_at = at.matrix("B")
Z_at = A_at.dot(x_at + B_at.dot(c_at + d_at))
fgraph = FunctionGraph(outputs=[Z_at], clone=False)
assert isinstance(fgraph.outputs[0].owner.op, Dot)
def distributes(in_lv, out_lv):
return lall(
# lhs == A * (x + b)
eq(
etuple(_dot, var("A"), etuple(at.add, var("x"), var("b"))),
in_lv,
),
# rhs == A * x + A * b
eq(
etuple(
at.add,
etuple(_dot, var("A"), var("x")),
etuple(_dot, var("A"), var("b")),
),
out_lv,
),
)
distribute_opt = EquilibriumOptimizer([KanrenRelationSub(distributes)],
max_use_ratio=10)
fgraph_opt = optimize_graph(fgraph, custom_opt=distribute_opt)
(expr_opt, ) = fgraph_opt.outputs
assert expr_opt.owner.op == at.add
assert isinstance(expr_opt.owner.inputs[0].owner.op, Dot)
assert fgraph_opt.inputs[0] is A_at
assert expr_opt.owner.inputs[0].owner.inputs[0].name == "A"
assert expr_opt.owner.inputs[1].owner.op == at.add
assert isinstance(expr_opt.owner.inputs[1].owner.inputs[0].owner.op, Dot)
assert isinstance(expr_opt.owner.inputs[1].owner.inputs[1].owner.op, Dot)
def test_infer_shape(self):
# test infer shape does not need to against inline case
# since the Op is remove during optimization phase
x = matrix("x")
y = matrix("y")
o1 = x + y
o2 = x * y
op_graph = OpFromGraph([x, y], [o1, o2])
q = matrix("q")
p = matrix("p")
self._compile_and_check(
[q, p],
op_graph(q, p),
[
np.ones([3, 4], dtype=config.floatX),
np.ones([3, 4], dtype=config.floatX),
],
OpFromGraph,
)
# Make sure `OpFromGraph.infer_shape` can handle objects without a
# shape
x = MyVariable("x")
y = matrix("y")
z = as_tensor([1, 2])
op_graph = OpFromGraph([x, y, z], [x, y])
op_var = op_graph(x, y, z)
fg = FunctionGraph(outputs=[op_var[1]], clone=False)
opt_res = optimize_graph(fg, custom_opt=ShapeOptimizer())
assert opt_res.shape_feature.shape_of[x] is None
assert opt_res.shape_feature.shape_of[z][0].data == 2