def test_toposort2(): g0 = Graph() g0.output = Constant(33) g1 = Graph() in0 = Constant(g0) in1 = Constant(1) v1 = Apply([in0, in1], g1) v2 = Apply([in0, v1, in1], g1) g1.output = v2 order = list(toposort(g1.return_)) _check_toposort(order, g1.return_, succ_incoming)
def _opt_fancy_transpose(optimizer, node, equiv): if equiv[V].value == (1, 0): x = equiv[X] ct = Constant(GraphCosmeticPrimitive(f'T', on_edge=True)) with About(node.debug, 'cosmetic'): return Apply([ct, x], node.graph) else: return node
def test_dfs_graphs(): g0 = Graph() in0 = Constant(g0) in1 = Constant(1) g0.return_ = in1 value = Apply([in0], Graph()) assert set(dfs(value)) == {value, in0} assert set(dfs(value, follow_graph=True)) == {value, in0, in1}
def _opt_fancy_array_map(optimizer, node, equiv): xs = equiv[Xs] v = equiv[V] if v.is_constant_graph(): return node name = short_labeler.label(v) ct = Constant(GraphCosmeticPrimitive(f'[{name}]')) with About(node.debug, 'cosmetic'): return Apply([ct, *xs], node.graph)
def test_toposort(): g0 = Graph() g0.output = Constant(1) g1 = Graph() in0 = Constant(g0) value = Apply([in0], g1) g1.output = value order = list(toposort(g1.return_)) _check_toposort(order, g1.return_, succ_incoming)
def test_helpers(): g = Graph() cg = Constant(g) assert cg.is_constant(Graph) assert cg.is_constant_graph() one = Constant(1) assert one.is_constant() assert one.is_constant(int) assert not one.is_constant(str) assert not one.is_constant_graph() a = Apply([cg, one], g) assert a.is_apply() p = Parameter(g) assert p.is_parameter() s = Special(1234, g) assert s.is_special() assert s.is_special(int) assert not s.is_special(str)
def _opt_fancy_sum(optimizer, node, equiv): x = equiv[X] shp = equiv[V].value ct = Constant(GraphCosmeticPrimitive(f'sum {"x".join(map(str, shp))}')) with About(node.debug, 'cosmetic'): return Apply([ct, x], node.graph)
def _opt_fancy_array_to_scalar(optimizer, node, equiv): x = equiv[X] ct = Constant(GraphCosmeticPrimitive(f'to_scalar', on_edge=True)) with About(node.debug, 'cosmetic'): return Apply([ct, x], node.graph)
def _opt_fancy_distribute(optimizer, node, equiv): x = equiv[X] v = equiv[V] ct = Constant(GraphCosmeticPrimitive(f'shape→{v.value}', on_edge=True)) with About(node.debug, 'cosmetic'): return Apply([ct, x], node.graph)
def _opt_fancy_unsafe_static_cast(optimizer, node, equiv): x = equiv[X] ct = Constant(GraphCosmeticPrimitive(f'cast', on_edge=True)) with About(node.debug, 'cosmetic'): return Apply([ct, x], node.graph)
def _opt_fancy_make_tuple(optimizer, node, equiv): xs = equiv[Xs] ct = Constant(GraphCosmeticPrimitive('(...)')) with About(node.debug, 'cosmetic'): return Apply([ct, *xs], node.graph)
def _opt_fancy_scalar_to_array(optimizer, node, equiv): x = equiv[X] ct = Constant(GraphCosmeticPrimitive(f"to_array", on_edge=True)) with About(node.debug, "cosmetic"): return Apply([ct, x], node.graph)
def _opt_fancy_tagged(optimizer, node, equiv): x = equiv[X] v = equiv[V] ct = Constant(GraphCosmeticPrimitive(f"@{v.value}", on_edge=True)) with About(node.debug, "cosmetic"): return Apply([ct, x], node.graph)
def test_dfs(): in0 = Constant(0) in1 = Constant(1) value = Apply([in0, in1], Graph()) assert next(dfs(value)) == value assert set(dfs(value)) == {value, in0, in1}