Removes a stack->unstack pattern from in_graph_def in a returned graph. Args: in_graph_def: Graph def to use as input. Returns: Simplified tuple (graph_def, changed_something) where changed_something is true if anything was done.
(in_graph_def)
| 1038 | |
| 1039 | # TODO(aselle): This should be converted to grappler in the future. |
| 1040 | def _remove_one_redundant_stack_unstack(in_graph_def): |
| 1041 | """Removes a stack->unstack pattern from in_graph_def in a returned graph. |
| 1042 | |
| 1043 | Args: |
| 1044 | in_graph_def: Graph def to use as input. |
| 1045 | Returns: |
| 1046 | Simplified tuple (graph_def, changed_something) where changed_something |
| 1047 | is true if anything was done. |
| 1048 | """ |
| 1049 | name_to_input_name, name_to_node, name_to_seq_num = _extract_graph_summary( |
| 1050 | in_graph_def) |
| 1051 | del name_to_seq_num |
| 1052 | |
| 1053 | # TODO(aselle): Make this not hardcoded. |
| 1054 | do_generic_pack_unpack = True |
| 1055 | |
| 1056 | out = _graph_pb2.GraphDef() |
| 1057 | out.library.CopyFrom(in_graph_def.library) |
| 1058 | out.versions.CopyFrom(in_graph_def.versions) |
| 1059 | for n in in_graph_def.node: |
| 1060 | node_name = _tensor_name_base(n.name) |
| 1061 | if not node_name.startswith("OpHintStack") and not n.op.startswith("Pack"): |
| 1062 | continue |
| 1063 | next_to_visit = [node_name] |
| 1064 | visited = set() |
| 1065 | |
| 1066 | unpack_nodes = set() |
| 1067 | pack_node = node_name |
| 1068 | |
| 1069 | # Find a pattern of unstack connected to a stack (with identities |
| 1070 | # in between. |
| 1071 | matches_pattern = True |
| 1072 | is_hint_created_stack = False |
| 1073 | while next_to_visit: |
| 1074 | current_node_name = next_to_visit[0] |
| 1075 | visited.add(current_node_name) |
| 1076 | del next_to_visit[0] |
| 1077 | node = name_to_node[current_node_name] |
| 1078 | is_op_hint_stack = node.name.startswith("OpHintStack") |
| 1079 | is_op_hint_unstack = node.name.startswith("OpHintUnstack") |
| 1080 | if (node.op == "Identity" or is_op_hint_stack |
| 1081 | or (do_generic_pack_unpack and node.op == "Pack")): |
| 1082 | is_hint_created_stack |= is_op_hint_stack |
| 1083 | next_to_visit += [ |
| 1084 | input_node for input_node in name_to_input_name[current_node_name] |
| 1085 | if input_node not in visited |
| 1086 | ] |
| 1087 | elif (is_op_hint_unstack |
| 1088 | or (do_generic_pack_unpack and node.op == "Unpack")): |
| 1089 | unpack_nodes.add(node.name) |
| 1090 | is_hint_created_stack &= is_op_hint_unstack |
| 1091 | else: |
| 1092 | matches_pattern = False |
| 1093 | break |
| 1094 | visited.add(node.name) |
| 1095 | |
| 1096 | if matches_pattern and len(unpack_nodes) == 1: |
| 1097 | pack_node = node_name |
no test coverage detected