def test_patternsub_values_eq_approx(out_pattern, tracks): # PatternSub would fail when `values_eq_approx` and `get_nodes` were specified x = MyVariable("x") e = op1(x) fg = FunctionGraph([x], [e], clone=False) opt = EquilibriumOptimizer( [ PatternSub( (op1, "x"), out_pattern, tracks=[op1] if tracks else (), get_nodes=(lambda fgraph, node: [node]) if tracks else None, values_eq_approx=values_eq_approx_always_true, ) ], max_use_ratio=1, ) opt.optimize(fg) output = fg.outputs[0] if isinstance(out_pattern, tuple): assert output.owner.op == op2 assert output.tag.values_eq_approx is values_eq_approx_always_true elif out_pattern == "x": assert output is x assert output.tag.values_eq_approx is values_eq_approx_always_true else: # The replacement types do not match, so the substitution should've # failed assert output is e
def test_Subtensor_lift_restrictions(): rng = shared(np.random.RandomState(1233532), borrow=False) std = vector("std") std.tag.test_value = np.array([1e-5, 2e-5, 3e-5], dtype=config.floatX) x = normal(aet.arange(2), aet.ones(2), rng=rng) y = x[1] # The non-`Subtensor` client depends on the RNG state, so we can't perform # the lift z = x - y fg = FunctionGraph([rng], [z], clone=False) _ = EquilibriumOptimizer([local_subtensor_rv_lift], max_use_ratio=100).apply(fg) subtensor_node = fg.outputs[0].owner.inputs[1].owner.inputs[0].owner assert subtensor_node == y.owner assert isinstance(subtensor_node.op, Subtensor) assert subtensor_node.inputs[0].owner.op == normal # The non-`Subtensor` client doesn't depend on the RNG state, so we can # perform the lift z = aet.ones(x.shape) - x[1] fg = FunctionGraph([rng], [z], clone=False) EquilibriumOptimizer([local_subtensor_rv_lift], max_use_ratio=100).apply(fg) rv_node = fg.outputs[0].owner.inputs[1].owner.inputs[0].owner assert rv_node.op == normal assert isinstance(rv_node.inputs[-1].owner.op, Subtensor) assert isinstance(rv_node.inputs[-2].owner.op, Subtensor)
def test_Dimshuffle_lift_restrictions(): rng = shared(np.random.RandomState(1233532), borrow=False) x = normal(aet.arange(2).reshape((2, )), 100, size=(2, 2, 2), rng=rng) y = x.dimshuffle(1, 0, 2) # The non-`Dimshuffle` client depends on the RNG state, so we can't # perform the lift z = x - y fg = FunctionGraph([rng], [z], clone=False) _ = EquilibriumOptimizer([local_dimshuffle_rv_lift], max_use_ratio=100).apply(fg) dimshuffle_node = fg.outputs[0].owner.inputs[1].owner assert dimshuffle_node == y.owner assert isinstance(dimshuffle_node.op, DimShuffle) assert dimshuffle_node.inputs[0].owner.op == normal # The non-`Dimshuffle` client doesn't depend on the RNG state, so we can # perform the lift z = aet.ones(x.shape) - y fg = FunctionGraph([rng], [z], clone=False) EquilibriumOptimizer([local_dimshuffle_rv_lift], max_use_ratio=100).apply(fg) rv_node = fg.outputs[0].owner.inputs[1].owner assert rv_node.op == normal assert isinstance(rv_node.inputs[-1].owner.op, DimShuffle) assert isinstance(rv_node.inputs[-2].owner.op, DimShuffle)
def test_patternsub_invalid_dtype(out_pattern): # PatternSub would wrongly return output of different dtype as the original node x = MyVariable("x") e = op_cast_type2(x) fg = FunctionGraph([x], [e]) opt = EquilibriumOptimizer( [PatternSub( (op_cast_type2, "x"), out_pattern, )], max_use_ratio=1, ) opt.optimize(fg) assert fg.apply_nodes.pop().op == op_cast_type2
def test_1(self): x, y, z = map(MyVariable, "xyz") e = op3(op4(x, y)) g = FunctionGraph([x, y, z], [e]) # print g opt = EquilibriumOptimizer( [ PatternSub((op1, "x", "y"), (op2, "x", "y")), PatternSub((op4, "x", "y"), (op1, "x", "y")), PatternSub((op3, (op2, "x", "y")), (op4, "x", "y")), ], max_use_ratio=10, ) opt.optimize(g) # print g assert str(g) == "FunctionGraph(Op2(x, y))"
def apply_local_opt_to_rv(opt, op_fn, dist_op, dist_params, size, rng): dist_params_aet = [] for p in dist_params: p_aet = aet.as_tensor(p).type() p_aet.tag.test_value = p dist_params_aet.append(p_aet) size_aet = [] for s in size: s_aet = iscalar() s_aet.tag.test_value = s size_aet.append(s_aet) dist_st = op_fn(dist_op(*dist_params_aet, size=size_aet, rng=rng)) f_inputs = [ p for p in dist_params_aet + size_aet if not isinstance(p, (slice, Constant)) ] mode = Mode("py", EquilibriumOptimizer([opt], max_use_ratio=100)) f_opt = function( f_inputs, dist_st, mode=mode, ) (new_out, ) = f_opt.maker.fgraph.outputs return new_out, f_inputs, dist_st, f_opt
def test_1(self): x, y, z = map(MyVariable, "xyz") # TODO FIXME: These `Op`s don't have matching/consistent `__prop__`s # and `__init__`s, so they can't be `etuplized` correctly e = op3(op4(x, y)) g = FunctionGraph([x, y, z], [e]) # print g opt = EquilibriumOptimizer( [ PatternSub((op1, "x", "y"), (op2, "x", "y")), PatternSub((op4, "x", "y"), (op1, "x", "y")), PatternSub((op3, (op2, "x", "y")), (op4, "x", "y")), ], max_use_ratio=10, ) opt.optimize(g) # print g assert str(g) == "FunctionGraph(Op2(x, y))"
def test_DimShuffle_lift(ds_order, lifted, dist_op, dist_params, size, rtol): rng = shared(np.random.RandomState(1233532), borrow=False) dist_params_aet = [] for p in dist_params: p_aet = aet.as_tensor(p).type() p_aet.tag.test_value = p dist_params_aet.append(p_aet) size_aet = [] for s in size: s_aet = iscalar() s_aet.tag.test_value = s size_aet.append(s_aet) dist_st = dist_op(*dist_params_aet, size=size_aet, rng=rng).dimshuffle(ds_order) f_inputs = [ p for p in dist_params_aet + size_aet if not isinstance(p, (slice, Constant)) ] mode = Mode( "py", EquilibriumOptimizer([local_dimshuffle_rv_lift], max_use_ratio=100) ) f_opt = function( f_inputs, dist_st, mode=mode, ) (new_out,) = f_opt.maker.fgraph.outputs if lifted: assert new_out.owner.op == dist_op assert all( isinstance(i.owner.op, DimShuffle) for i in new_out.owner.inputs[3:] if i.owner ) else: assert isinstance(new_out.owner.op, DimShuffle) return f_base = function( f_inputs, dist_st, mode=no_mode, ) arg_values = [p.get_test_value() for p in f_inputs] res_base = f_base(*arg_values) res_opt = f_opt(*arg_values) np.testing.assert_allclose(res_base, res_opt, rtol=rtol)
def test_Subtensor_lift_restrictions(): rng = shared(np.random.default_rng(1233532), borrow=False) std = vector("std") std.tag.test_value = np.array([1e-5, 2e-5, 3e-5], dtype=config.floatX) x = normal(at.arange(2), at.ones(2), rng=rng) y = x[1] # The non-`Subtensor` client depends on the RNG state, so we can't perform # the lift z = x - y fg = FunctionGraph([rng], [z], clone=False) _ = EquilibriumOptimizer([local_subtensor_rv_lift], max_use_ratio=100).apply(fg) subtensor_node = fg.outputs[0].owner.inputs[1].owner.inputs[0].owner assert subtensor_node == y.owner assert isinstance(subtensor_node.op, Subtensor) assert subtensor_node.inputs[0].owner.op == normal z = at.ones(x.shape) - x[1] # We add `x` as an output to make sure that `is_rv_used_in_graph` handles # `"output"` "nodes" correctly. fg = FunctionGraph([rng], [z, x], clone=False) EquilibriumOptimizer([local_subtensor_rv_lift], max_use_ratio=100).apply(fg) assert fg.outputs[0] == z assert fg.outputs[1] == x # The non-`Subtensor` client doesn't depend on the RNG state, so we can # perform the lift fg = FunctionGraph([rng], [z], clone=False) EquilibriumOptimizer([local_subtensor_rv_lift], max_use_ratio=100).apply(fg) rv_node = fg.outputs[0].owner.inputs[1].owner.inputs[0].owner assert rv_node.op == normal assert isinstance(rv_node.inputs[-1].owner.op, Subtensor) assert isinstance(rv_node.inputs[-2].owner.op, Subtensor)
def test_Dimshuffle_lift_restrictions(): rng = shared(np.random.default_rng(1233532), borrow=False) x = normal(at.arange(2).reshape((2, )), 100, size=(2, 2, 2), rng=rng) y = x.dimshuffle(1, 0, 2) # The non-`Dimshuffle` client depends on the RNG state, so we can't # perform the lift z = x - y fg = FunctionGraph([rng], [z, y], clone=False) _ = EquilibriumOptimizer([local_dimshuffle_rv_lift], max_use_ratio=100).apply(fg) dimshuffle_node = fg.outputs[0].owner.inputs[1].owner assert dimshuffle_node == y.owner assert isinstance(dimshuffle_node.op, DimShuffle) assert dimshuffle_node.inputs[0].owner.op == normal z = at.ones(x.shape) - y # We add `x` as an output to make sure that `is_rv_used_in_graph` handles # `"output"` "nodes" correctly. fg = FunctionGraph([rng], [z, x], clone=False) EquilibriumOptimizer([local_dimshuffle_rv_lift], max_use_ratio=100).apply(fg) assert fg.outputs[0] == z assert fg.outputs[1] == x # The non-`Dimshuffle` client doesn't depend on the RNG state, so we can # perform the lift fg = FunctionGraph([rng], [z], clone=False) EquilibriumOptimizer([local_dimshuffle_rv_lift], max_use_ratio=100).apply(fg) rv_node = fg.outputs[0].owner.inputs[1].owner assert rv_node.op == normal assert isinstance(rv_node.inputs[-1].owner.op, DimShuffle) assert isinstance(rv_node.inputs[-2].owner.op, DimShuffle)
def test_low_use_ratio(self): x, y, z = map(MyVariable, "xyz") e = op3(op4(x, y)) g = FunctionGraph([x, y, z], [e]) # print 'before', g # display pesky warnings along with stdout # also silence logger for 'aesara.graph.opt' _logger = logging.getLogger("aesara.graph.opt") oldlevel = _logger.level _logger.setLevel(logging.CRITICAL) try: opt = EquilibriumOptimizer( [ PatternSub((op1, "x", "y"), (op2, "x", "y")), PatternSub((op4, "x", "y"), (op1, "x", "y")), PatternSub((op3, (op2, "x", "y")), (op4, "x", "y")), ], max_use_ratio=1.0 / len(g.apply_nodes), ) # each opt can only be applied once opt.optimize(g) finally: _logger.setLevel(oldlevel) # print 'after', g assert str(g) == "FunctionGraph(Op1(x, 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_Subtensor_lift(indices, lifted, dist_op, dist_params, size): rng = shared(np.random.RandomState(1233532), borrow=False) dist_params_aet = [] for p in dist_params: p_aet = aet.as_tensor(p).type() p_aet.tag.test_value = p dist_params_aet.append(p_aet) size_aet = [] for s in size: s_aet = iscalar() s_aet.tag.test_value = s size_aet.append(s_aet) from aesara.tensor.subtensor import as_index_constant indices_aet = () for i in indices: i_aet = as_index_constant(i) if not isinstance(i_aet, slice): i_aet.tag.test_value = i indices_aet += (i_aet,) dist_st = dist_op(*dist_params_aet, size=size_aet, rng=rng)[indices_aet] f_inputs = [ p for p in dist_params_aet + size_aet + list(indices_aet) if not isinstance(p, (slice, Constant)) ] mode = Mode( "py", EquilibriumOptimizer([local_subtensor_rv_lift], max_use_ratio=100) ) f_opt = function( f_inputs, dist_st, mode=mode, ) (new_out,) = f_opt.maker.fgraph.outputs if lifted: assert isinstance(new_out.owner.op, RandomVariable) assert all( isinstance(i.owner.op, (AdvancedSubtensor, AdvancedSubtensor1, Subtensor)) for i in new_out.owner.inputs[3:] if i.owner ) else: assert isinstance( new_out.owner.op, (AdvancedSubtensor, AdvancedSubtensor1, Subtensor) ) return f_base = function( f_inputs, dist_st, mode=no_mode, ) arg_values = [p.get_test_value() for p in f_inputs] res_base = f_base(*arg_values) res_opt = f_opt(*arg_values) np.testing.assert_allclose(res_base, res_opt, rtol=1e-3)