(
self,
flat_graph,
nodes,
seen_nodes,
seen_modules,
parent,
module_stack,
module_id,
module_call_graph: Dict[str, ModuleCallSignature],
module: Optional[torch.nn.Module] = None,
)
| 306 | |
| 307 | class ModuleFrame: |
| 308 | def __init__( |
| 309 | self, |
| 310 | flat_graph, |
| 311 | nodes, |
| 312 | seen_nodes, |
| 313 | seen_modules, |
| 314 | parent, |
| 315 | module_stack, |
| 316 | module_id, |
| 317 | module_call_graph: Dict[str, ModuleCallSignature], |
| 318 | module: Optional[torch.nn.Module] = None, |
| 319 | ): |
| 320 | self.flat_graph = flat_graph |
| 321 | self.nodes = nodes |
| 322 | self.seen_nodes = seen_nodes |
| 323 | self.seen_modules = seen_modules |
| 324 | self.parent = parent |
| 325 | self.module_stack = module_stack |
| 326 | self.module_id = module_id |
| 327 | |
| 328 | self.module_call_graph = module_call_graph |
| 329 | self.verbose = False |
| 330 | |
| 331 | self.fqn = self.module_stack[-1] |
| 332 | if module is not None: |
| 333 | self.module = module |
| 334 | else: |
| 335 | self.module = InterpreterModule( |
| 336 | torch.fx.Graph(), module_call_graph.get(self.fqn) |
| 337 | ) |
| 338 | if self.module_id in self.seen_modules: |
| 339 | self.cached_graph_module = self.seen_modules[self.module_id] |
| 340 | else: |
| 341 | self.cached_graph_module = None |
| 342 | self.seen_modules[self.module_id] = self.module |
| 343 | |
| 344 | self.graph = self.module.graph |
| 345 | |
| 346 | # Mapping of nodes in the flat graph to nodes in this graph. |
| 347 | self.node_map: Dict[torch.fx.Node, torch.fx.Node] = {} |
| 348 | self.node_to_placeholder = {} |
| 349 | |
| 350 | self.parent_call_module: Optional[torch.fx.Node] = None |
| 351 | if parent is not None: |
| 352 | accessor = compute_accessor(parent.fqn, self.fqn) |
| 353 | parent.module.add_module( |
| 354 | accessor, |
| 355 | self.module |
| 356 | if self.cached_graph_module is None |
| 357 | else self.cached_graph_module, |
| 358 | ) |
| 359 | self.parent_call_module = parent.graph.call_module(accessor) |
| 360 | |
| 361 | signature = module_call_graph.get(self.fqn) |
| 362 | if signature is not None and self.parent is not None: |
| 363 | assert len(signature.in_spec.children_specs) == 2 |
| 364 | args_spec = signature.in_spec.children_specs[0] |
| 365 | kwargs_spec = signature.in_spec.children_specs[1] |
no test coverage detected