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}
