Converts an InputSpec or OutputSpec to its corresponding node in the graph.
(
exp_program: ExportedProgram, spec: InputSpec | OutputSpec
)
| 250 | |
| 251 | |
| 252 | def _spec_to_node( |
| 253 | exp_program: ExportedProgram, spec: InputSpec | OutputSpec |
| 254 | ) -> torch.fx.Node: |
| 255 | """ |
| 256 | Converts an InputSpec or OutputSpec to its corresponding node in the graph. |
| 257 | """ |
| 258 | # Extract the argument name from the spec |
| 259 | if hasattr(spec, "arg") and hasattr(spec.arg, "name"): |
| 260 | arg_name = spec.arg.name |
| 261 | else: |
| 262 | raise RuntimeError(f"Invalid spec format: {spec}") |
| 263 | |
| 264 | # Find the corresponding node in the graph |
| 265 | for node in exp_program.graph.nodes: |
| 266 | if node.name == arg_name: |
| 267 | return node |
| 268 | |
| 269 | raise RuntimeError(f"Could not find node with name '{arg_name}' in the graph") |
| 270 | |
| 271 | |
| 272 | def create_mutable_buffer( |
no outgoing calls
no test coverage detected