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_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_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_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}