Find all children hints. For a given OpHint, we find all children hints inside it, we also copy all the nodes inside function defs (if applicable) to the original graph_def, they are returned in a list as well. Args: call: Parent OpHint that contains children ophints. graph_def: Or
(call, graph_def)
| 822 | |
| 823 | |
| 824 | def _find_children_hints(call, graph_def): |
| 825 | """Find all children hints. |
| 826 | |
| 827 | For a given OpHint, we find all children hints inside it, we also copy all the |
| 828 | nodes inside function defs (if applicable) to the original graph_def, they are |
| 829 | returned in a list as well. |
| 830 | |
| 831 | Args: |
| 832 | call: Parent OpHint that contains children ophints. |
| 833 | graph_def: Original graph def. |
| 834 | |
| 835 | Returns: |
| 836 | Ordered children hints inside the parent ophint; new graph def that contains |
| 837 | nodes inside function defs (if applicable); nodes inside function defs. |
| 838 | """ |
| 839 | name_to_input_name, _, _ = _extract_graph_summary(graph_def) |
| 840 | input_names, output_names = call.flattened_inputs_and_outputs() |
| 841 | |
| 842 | reachable_by_input = _bfs_for_reachable_nodes(input_names, name_to_input_name) |
| 843 | reachable_by_output = _bfs_for_reachable_nodes(output_names, |
| 844 | name_to_input_name) |
| 845 | output_nodes_set = set(output_names) |
| 846 | children_hints = [] |
| 847 | out = _graph_pb2.GraphDef() |
| 848 | out.library.CopyFrom(graph_def.library) |
| 849 | out.versions.CopyFrom(graph_def.versions) |
| 850 | function_def_nodes = set() |
| 851 | for node in graph_def.node: |
| 852 | out.node.extend([_copy.deepcopy(node)]) |
| 853 | n = _tensor_name_base(node.name) |
| 854 | if n in reachable_by_output: |
| 855 | if n not in reachable_by_input and n not in output_nodes_set: |
| 856 | # special handle for while loop function def. |
| 857 | if node.op == "While" or node.op == "StatelessWhile": |
| 858 | body_name = node.attr["body"].func.name |
| 859 | inputs_outside_loop = node.input |
| 860 | for function_def in graph_def.library.function: |
| 861 | if function_def.signature.name == body_name: |
| 862 | function_inputs = function_def.signature.input_arg |
| 863 | assert len(inputs_outside_loop) == len(function_inputs) |
| 864 | nodes_mapping = {} |
| 865 | for i, function_input in enumerate(function_inputs): |
| 866 | nodes_mapping[function_input.name] = inputs_outside_loop[i] |
| 867 | # TODO(b/123050804): Consider use grappler. |
| 868 | (children_hints_in_loop, |
| 869 | new_nodes) = _find_children_hints_in_while_loop( |
| 870 | function_def, nodes_mapping) |
| 871 | function_def_nodes.update([x.name for x in new_nodes]) |
| 872 | children_hints.extend(children_hints_in_loop) |
| 873 | out.node.extend(new_nodes) |
| 874 | |
| 875 | return children_hints, out, function_def_nodes |
| 876 | |
| 877 | |
| 878 | def _tensor_name_base(full_tensor_name): |
no test coverage detected