def label(self, node, force=None, fn_label=None):
"""Label a node."""
if isinstance(node, DebugInfo):
return self.name(node, True if force is None else force)
elif isinstance(node, Graph):
return self.name(node.debug, True if force is None else force)
elif is_constant_graph(node):
return self.name(node.value.debug,
True if force is None else force)
elif is_constant(node):
v = node.value
if isinstance(v, (int, float, str)):
return repr(v)
elif isinstance(v, Primitive):
return v.name
else:
class_name = v.__class__.__name__
return f'{self.label(node.debug, True)}:{class_name}'
elif is_parameter(node):
return self.label(node.debug, True)
else:
lbl = ''
if self.function_in_node:
if fn_label is None:
fn_label = self.const_fn(node)
if fn_label:
lbl = fn_label
name = self.name(node, force)
if name:
if lbl:
lbl += f'→{name}'
else:
lbl = name
return lbl or '·'
def process_edges(self, edges):
"""Create edges."""
for edge in edges:
src, lbl, dest = edge
if is_constant(dest) and self.duplicate_constants:
self.follow(dest)
cid = self.fresh_id()
self.cynode(id=cid,
parent=src.graph,
label=self.label(dest),
classes=self.class_gen(dest, 'constant'),
node=dest)
self.cyedge(src_id=src, dest_id=cid, label=lbl)
elif self.duplicate_free_variables and \
src.graph and dest.graph and \
src.graph is not dest.graph:
self.pool.add(dest)
cid = self.fresh_id()
self.cynode(id=cid,
parent=src.graph,
label=self.name(dest),
classes=self.class_gen(dest, 'freevar'),
node=dest)
self.cyedge(src_id=src, dest_id=cid, label=lbl)
self.cyedge(src_id=cid, dest_id=dest, label=(lbl, 'link-edge'))
else:
self.pool.add(dest)
self.cyedge(src_id=src, dest_id=dest, label=lbl)
def process_node_getitem(self, node, g, cl):
"""Create node and edges for `x[ct]`."""
idx = node.inputs[2]
if self.function_in_node and is_constant(idx):
lbl = self.label(node, '')
self.cynode(id=node, label=lbl, parent=g, classes=cl)
self.process_edges([(node, (f'[{idx.value}]', 'fn-edge'),
node.inputs[1])])
else:
self.process_node_generic(node, g, cl)
def const_fn(self, node):
"""
Return name of function, if constant.
Given an `Apply` node of a constant function, return the
name of that function, otherwise return None.
"""
fn = node.inputs[0] if node.inputs else None
if fn and is_constant(fn):
return self.label(fn, False)
else:
return None
def test_helpers():
g = Graph()
cg = Constant(g)
assert is_constant_graph(cg)
one = Constant(1)
assert not is_constant_graph(one)
assert is_constant(one)
a = Apply([cg, one], g)
assert is_apply(a)
p = Parameter(g)
assert is_parameter(p)
def class_gen(self, node, cl=None):
"""Generate the class name for this node."""
g = node.graph
if cl is not None:
pass
elif node in self.returns:
cl = 'output'
elif g and node in g.parameters:
cl = 'input'
elif is_constant(node):
cl = 'constant'
else:
cl = 'intermediate'
if self._class_gen:
return self._class_gen(node, cl)
else:
return cl
def process_node_generic(self, node, g, cl):
"""Create node and edges for a node."""
lbl = self.label(node)
self.cynode(id=node, label=lbl, parent=g, classes=cl)
fn = node.inputs[0] if node.inputs else None
if fn and is_constant_graph(fn):
self.graphs.add(fn.value)
edges = []
if fn and not (is_constant(fn) and self.function_in_node):
edges.append((node, 'F', fn))
edges += [(node, i + 1, inp)
for i, inp in enumerate(node.inputs[1:]) or []]
self.process_edges(edges)