Normalize a serialized ExportedProgram, so that different eager program which shares the same semantics can get a single representation on disk. This function canonicalizes an ExportedProgram by: 1. Sorting nodes in topological order. 2. Rename nodes to have unique names.
(ep: ExportedProgram)
| 2681 | |
| 2682 | |
| 2683 | def canonicalize(ep: ExportedProgram) -> ExportedProgram: |
| 2684 | """ |
| 2685 | Normalize a serialized ExportedProgram, so that different eager program which |
| 2686 | shares the same semantics can get a single representation on disk. |
| 2687 | |
| 2688 | This function canonicalizes an ExportedProgram by: |
| 2689 | |
| 2690 | 1. Sorting nodes in topological order. |
| 2691 | 2. Rename nodes to have unique names. |
| 2692 | 3. Remove unstable fields. |
| 2693 | 4. Aggregate the above program fields. |
| 2694 | 5. Recurse in subgraphs. |
| 2695 | |
| 2696 | Args: |
| 2697 | ep (ExportedProgram): The ExportedProgram to canonicalize. |
| 2698 | |
| 2699 | Returns: |
| 2700 | ExportedProgram: The canonicalized exported program. |
| 2701 | """ |
| 2702 | ep = copy.deepcopy(ep) |
| 2703 | |
| 2704 | opset_version = dict(sorted(ep.opset_version.items(), key=operator.itemgetter(0))) |
| 2705 | range_constraints = dict( |
| 2706 | sorted(ep.range_constraints.items(), key=operator.itemgetter(0)) |
| 2707 | ) |
| 2708 | module_call_graph = sorted(ep.graph_module.module_call_graph, key=lambda x: x.fqn) |
| 2709 | signature = ep.graph_module.signature |
| 2710 | graph = ep.graph_module.graph |
| 2711 | |
| 2712 | assert len(graph.inputs) == len(signature.input_specs) |
| 2713 | assert len(graph.outputs) == len(signature.output_specs) |
| 2714 | |
| 2715 | def rank_input(inp) -> Tuple[int, Optional[str], int]: |
| 2716 | idx, (arg, spec) = inp |
| 2717 | assert isinstance(spec, InputSpec) |
| 2718 | if spec.type == "user_input": |
| 2719 | return 5, None, idx |
| 2720 | elif spec.type == "parameter": |
| 2721 | return 1, spec.parameter.parameter_name, idx |
| 2722 | elif spec.type == "buffer": |
| 2723 | return 2, spec.buffer.buffer_name, idx |
| 2724 | elif spec.type == "tensor_constant": |
| 2725 | return 3, spec.tensor_constant.tensor_constant_name, idx |
| 2726 | elif spec.type == "custom_obj": |
| 2727 | return 4, spec.custom_obj.custom_obj_name, idx |
| 2728 | elif spec.type == "token": |
| 2729 | return 0, None, idx |
| 2730 | elif spec.type == "constant_input": |
| 2731 | return 6, spec.constant_input.name, idx |
| 2732 | else: |
| 2733 | raise AssertionError(f"Unknown input type: {spec}") |
| 2734 | |
| 2735 | def rank_output(out) -> Tuple[int, Optional[str], int]: |
| 2736 | idx, (arg, spec) = out |
| 2737 | assert isinstance(spec, OutputSpec) |
| 2738 | if spec.type == "user_output": |
| 2739 | return 3, None, idx |
| 2740 | elif spec.type == "loss_output": |
no test coverage detected