(
node: Node,
base_mod_env: Dict[str, Node],
base_mod_attrs: Dict[str, torch.fx.graph_module.GraphModule],
)
| 141 | """ |
| 142 | |
| 143 | def construct_graph( |
| 144 | node: Node, |
| 145 | base_mod_env: Dict[str, Node], |
| 146 | base_mod_attrs: Dict[str, torch.fx.graph_module.GraphModule], |
| 147 | ): |
| 148 | if node.op == "placeholder": |
| 149 | default_value = ( |
| 150 | node.args[0] if len(node.args) > 0 else inspect.Signature.empty |
| 151 | ) |
| 152 | base_mod_env[node.name] = base_mod_graph.placeholder( |
| 153 | node.target, type_expr=node.type, default_value=default_value |
| 154 | ) |
| 155 | base_mod_env[node.name].meta = node.meta.copy() |
| 156 | elif node.op == "get_attr": |
| 157 | base_mod_env[node.name] = base_mod_graph.get_attr(node.target) |
| 158 | base_mod_env[node.name].meta = node.meta.copy() |
| 159 | attr_val = m |
| 160 | for atom in node.target.split("."): # type: ignore[union-attr] |
| 161 | if not hasattr(attr_val, atom): |
| 162 | raise AttributeError(f"Node target {node.target} not found!") |
| 163 | attr_val = getattr(attr_val, atom) |
| 164 | base_mod_attrs[node.target] = attr_val # type: ignore[index] |
| 165 | return base_mod_env, base_mod_attrs |
| 166 | |
| 167 | partitions: Dict[str, Partition] = {} |
| 168 | orig_nodes: Dict[str, Node] = {} |
no test coverage detected
searching dependent graphs…