def test_parse_tensor(self):
# Zero-rank tensor
attr = attr_value.AttrValue()
attr.tensor.version_number = 1
attr.tensor.dtype = types.DataType.DT_INT32
t = parse.parse_attr(attr)
self.assertTrue(isinstance(t, mil_types.int32))
self.assertEqual(0, t.val)
# Non-zero rank
attr = attr_value.AttrValue()
attr.tensor.version_number = 1
attr.tensor.dtype = types.DataType.DT_INT32
shaped_attr = self._attr_with_shape([(1, "outer"), (2, "middle"),
(3, "inner")])
attr.tensor.tensor_shape.dim.extend(shaped_attr.shape.dim)
attr.tensor.int_val.extend([55, 56, 57])
t = parse.parse_attr(attr)
self.assertEqual([55, 56, 57], t.val.tolist())
self.assertEqual("tensor", mil_types.get_type_info(t).name)
# Note that the result of t.get_primitive() is a function that returns a type
# rather than an instance of that type as it is when the tensor has rank zero.
self.assertTrue(isinstance(t.get_primitive()(), mil_types.int32))
self.assertEqual((1, 2, 3), t.get_shape())
def visit(self, graph, node, nodename_prefix=""):
if node.name in self.visited_memo:
return self
# For printing datatype, breaks type
if node.attr.get("symbolic_datatype", None) is not None:
dtype = str(types.get_type_info(node.attr["symbolic_datatype"]))
elif node.datatype is not None:
dtype = str(types.get_type_info(node.datatype))
else:
dtype = "Unknown"
label = ""
if self.alternate_labeller is not None:
label = self.alternate_labeller(node)
else:
if len(node.outputs) == 0:
label = "\\n{" + node.name + "}"
if "Placeholder" in node.op:
label = "\\n{" + node.name + "}"
if node.op == "while":
label = ("\\n{body: " + node.attr["body_function"] + " cond:" +
node.attr["cond_function"] + "}")
if node.op == "function":
label = "\\n{body: " + node.attr["function_name"] + "}"
if node.op == "function_entry":
label = "\\n{" + node.name + "}"
label = node.op + ":" + dtype + label
if node.name in self.highlights:
self.result.append('"' + nodename_prefix + node.name + '"' +
'[label="' + label +
'",fillcolor=%s,style=filled,fontcolor=%s]' % (
self.highlights[node.name],
"violetred" if node.attr.
get(self.annotation, False) else "black",
))
else:
self.result.append('"' + nodename_prefix + node.name + '"' +
'[label="' + label + '",fontcolor=%s]' %
("violetred" if node.attr.
get(self.annotation, False) else "black"))
for i in node.inputs:
input_name = i
edge = ('"' + nodename_prefix + input_name + '"' + " -> " + '"' +
nodename_prefix + node.name + '"')
innode = graph[input_name]
self.result.append(edge)
for i in node.control_inputs:
input_name = i
edge = ('"' + nodename_prefix + input_name + '"' + " -> " + '"' +
nodename_prefix + node.name + '"')
innode = graph[input_name]
edge = edge + " [style=dotted]"
self.result.append(edge)
self.visited_memo[node.name] = 1
for i in node.inputs:
input_name = i
if input_name[0] == "^":
input_name = input_name[1:]
assert input_name in graph
self.visit(graph, graph[input_name], nodename_prefix)
return self