(
tagged_graph_module: torch.fx.GraphModule,
owning_program: ExportedProgram,
field: str = "_in_target_subgraph",
)
| 146 | |
| 147 | |
| 148 | def extract_submodule_program( |
| 149 | tagged_graph_module: torch.fx.GraphModule, |
| 150 | owning_program: ExportedProgram, |
| 151 | field: str = "_in_target_subgraph", |
| 152 | ) -> ExportedProgram: |
| 153 | tagged_graph_module_output_node = tagged_graph_module.graph.output_node() |
| 154 | |
| 155 | partitioner = FlagBasedPartitioner(field) |
| 156 | partition_result = partitioner.partition(owning_program) |
| 157 | |
| 158 | tag, delegation_spec = next(iter(partition_result.partition_tags.items())) |
| 159 | node_list = _get_node_list_with_same_tag(tagged_graph_module, tag, owning_program) |
| 160 | |
| 161 | replace_ctx = tagged_graph_module._set_replace_hook( |
| 162 | owning_program.graph_signature.get_replace_hook() |
| 163 | ) |
| 164 | with replace_ctx: |
| 165 | submodule, call_module_node = create_submodule_from_nodes( |
| 166 | tagged_graph_module, node_list, tag |
| 167 | ) |
| 168 | |
| 169 | submodule_output_node = submodule.graph.output_node() |
| 170 | # Copy the output node meta from the original output node, because |
| 171 | # create_submodule_from_nodes doesn't cover the meta field |
| 172 | submodule_output_node.meta = tagged_graph_module_output_node.meta |
| 173 | |
| 174 | ( |
| 175 | submodule_program, |
| 176 | _, |
| 177 | _, |
| 178 | ) = create_exported_program_from_submodule( |
| 179 | submodule, |
| 180 | owning_program, |
| 181 | tag, |
| 182 | call_module_node, |
| 183 | False, |
| 184 | ) |
| 185 | |
| 186 | return submodule_program |
| 187 | |
| 188 | |
| 189 | class QuantizationMode(Enum): |
nothing calls this directly
no test coverage detected