def cosmetic_transformer(g):
"""Transform a graph so that it looks nicer.
The resulting graph is not a valid one to run, because it may contain nodes
with fake functions that only serve a cosmetic purpose.
"""
spec = (
_opt_distributed_constant,
_opt_fancy_make_tuple,
_opt_fancy_getitem,
_opt_fancy_resolve,
_opt_fancy_record_getitem,
_opt_fancy_array_map,
_opt_fancy_distribute,
_opt_fancy_transpose,
_opt_fancy_sum,
_opt_fancy_unsafe_static_cast,
_opt_fancy_scalar_to_array,
_opt_fancy_array_to_scalar,
_opt_fancy_hastag,
_opt_fancy_casttag,
_opt_fancy_tagged,
# careful=True
)
nmap = NodeMap()
for optim in spec:
nmap.register(getattr(optim, 'interest', None), optim)
optim = LocalPassOptimizer(nmap)
optim(g)
return g
def _check_opt(before, after, *opts, argspec=None, argspec_after=None):
nmap = NodeMap()
for opt in opts:
nmap.register(getattr(opt, 'interest', None), opt)
eq = LocalPassOptimizer(nmap)
_check_transform(before, after, eq, argspec, argspec_after)