Created
November 25, 2025 13:46
-
-
Save rec/a86a15be115b9e11444ea0aed2e1da3f to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| >>> Lint for torch/fx/passes/splitter_base.py: | |
| Error (MISSING_TYPE_LINTER) Missing return type for torch.fx.passes.splitter_base.NodeEvent.__init__ | |
| 142 | dep: Optional dependency, usually the node that caused the event. | |
| 143 | """ | |
| 144 | | |
| >>> 145 | def __init__( | |
| 146 | self, source: torch.fx.Node, desc: str, dep: Optional[torch.fx.Node] = None | |
| 147 | ): | |
| 148 | self.source = source | |
| Error (MISSING_TYPE_LINTER) Missing return type for torch.fx.passes.splitter_base.NodeEvent.to_str | |
| 149 | self.desc = desc | |
| 150 | self.dep = dep | |
| 151 | | |
| >>> 152 | def to_str(self): | |
| 153 | # source: The name of the subject of the event. | |
| 154 | # desc: description of the event, in the format of <event_type>|<explanation> | |
| 155 | # dep: The name of the cause of this event, which is another node, or # | |
| Error (MISSING_TYPE_LINTER) Missing return type for torch.fx.passes.splitter_base.NodeEventTracker.__init__ | |
| 163 | Tracks node events during the splitter execution. | |
| 164 | """ | |
| 165 | | |
| >>> 166 | def __init__(self, tracker_mode, dump_prefix): | |
| 167 | self.tracker_mode = tracker_mode | |
| 168 | self.dump_prefix = dump_prefix | |
| 169 | # list of events | |
| Error (MISSING_TYPE_LINTER) Missing return type for torch.fx.passes.splitter_base.NodeEventTracker.add | |
| 172 | self.node_events = {} | |
| 173 | self.writer = print | |
| 174 | | |
| >>> 175 | def add(self, node: torch.fx.Node, desc: str, dep: Optional[torch.fx.Node] = None): | |
| 176 | """ | |
| 177 | Add a new event to the tracker. | |
| 178 | """ | |
| Error (MISSING_TYPE_LINTER) Missing return type for torch.fx.passes.splitter_base.NodeEventTracker.print_node | |
| 182 | self.node_events[node.name] = [] | |
| 183 | self.node_events[node.name].append(len(self.events) - 1) | |
| 184 | | |
| >>> 185 | def print_node(self, node_name, recursive=False, tab="", writer=None): | |
| 186 | """ | |
| 187 | Print a node and its events. | |
| 188 | @param recursive: if True, print nodes that caused the events on this current node. | |
| Error (MISSING_TYPE_LINTER) Missing parameter type for torch.fx.passes.splitter_base.NodeEventTracker.print_node(node_name=) | |
| 182 | self.node_events[node.name] = [] | |
| 183 | self.node_events[node.name].append(len(self.events) - 1) | |
| 184 | | |
| >>> 185 | def print_node(self, node_name, recursive=False, tab="", writer=None): | |
| 186 | """ | |
| 187 | Print a node and its events. | |
| 188 | @param recursive: if True, print nodes that caused the events on this current node. | |
| Error (MISSING_TYPE_LINTER) Missing parameter type for torch.fx.passes.splitter_base.NodeEventTracker.print_node(recursive=) | |
| 182 | self.node_events[node.name] = [] | |
| 183 | self.node_events[node.name].append(len(self.events) - 1) | |
| 184 | | |
| >>> 185 | def print_node(self, node_name, recursive=False, tab="", writer=None): | |
| 186 | """ | |
| 187 | Print a node and its events. | |
| 188 | @param recursive: if True, print nodes that caused the events on this current node. | |
| Error (MISSING_TYPE_LINTER) Missing parameter type for torch.fx.passes.splitter_base.NodeEventTracker.print_node(tab=) | |
| 182 | self.node_events[node.name] = [] | |
| 183 | self.node_events[node.name].append(len(self.events) - 1) | |
| 184 | | |
| >>> 185 | def print_node(self, node_name, recursive=False, tab="", writer=None): | |
| 186 | """ | |
| 187 | Print a node and its events. | |
| 188 | @param recursive: if True, print nodes that caused the events on this current node. | |
| Error (MISSING_TYPE_LINTER) Missing parameter type for torch.fx.passes.splitter_base.NodeEventTracker.print_node(writer=) | |
| 182 | self.node_events[node.name] = [] | |
| 183 | self.node_events[node.name].append(len(self.events) - 1) | |
| 184 | | |
| >>> 185 | def print_node(self, node_name, recursive=False, tab="", writer=None): | |
| 186 | """ | |
| 187 | Print a node and its events. | |
| 188 | @param recursive: if True, print nodes that caused the events on this current node. | |
| Error (MISSING_TYPE_LINTER) Missing return type for torch.fx.passes.splitter_base.NodeEventTracker.to_dict | |
| 199 | event.dep.name, recursive=True, tab="| " + tab, writer=writer | |
| 200 | ) | |
| 201 | | |
| >>> 202 | def to_dict(self): | |
| 203 | """ | |
| 204 | Create dict dump on all events. | |
| 205 | """ | |
| Error (MISSING_TYPE_LINTER) Missing return type for torch.fx.passes.splitter_base.NodeEventTracker.print_all | |
| 211 | ret[name].append(event.to_str()) | |
| 212 | return ret | |
| 213 | | |
| >>> 214 | def print_all(self, writer=None): | |
| 215 | """ | |
| 216 | Print all nodes in a list. | |
| 217 | @param writer: function to write to file. If None, use print. | |
| Error (MISSING_TYPE_LINTER) Missing parameter type for torch.fx.passes.splitter_base.NodeEventTracker.print_all(writer=) | |
| 211 | ret[name].append(event.to_str()) | |
| 212 | return ret | |
| 213 | | |
| >>> 214 | def print_all(self, writer=None): | |
| 215 | """ | |
| 216 | Print all nodes in a list. | |
| 217 | @param writer: function to write to file. If None, use print. | |
| Error (MISSING_TYPE_LINTER) Missing return type for torch.fx.passes.splitter_base.NodeEventTracker.dump | |
| 222 | writer(f"Node: {name}:") | |
| 223 | self.print_node(name, recursive=False, tab=" ", writer=writer) | |
| 224 | | |
| >>> 225 | def dump(self): | |
| 226 | """ | |
| 227 | Function to be invoked at the end of the finder execution to printout tracked events specified by the mode. | |
| 228 | """ | |
| Error (MISSING_TYPE_LINTER) Missing return type for torch.fx.passes.splitter_base.NodeEventTracker.dump.writeln.fn | |
| 236 | payload_fn=lambda: json.dumps(self.to_dict()), | |
| 237 | ) | |
| 238 | | |
| >>> 239 | def writeln(f): | |
| 240 | def fn(x): | |
| 241 | return f.write(x + "\n") | |
| 242 | | |
| Error (MISSING_TYPE_LINTER) Missing parameter type for torch.fx.passes.splitter_base.NodeEventTracker.dump.writeln.fn(f=) | |
| 236 | payload_fn=lambda: json.dumps(self.to_dict()), | |
| 237 | ) | |
| 238 | | |
| >>> 239 | def writeln(f): | |
| 240 | def fn(x): | |
| 241 | return f.write(x + "\n") | |
| 242 | | |
| Error (MISSING_TYPE_LINTER) Missing return type for torch.fx.passes.splitter_base.NodeEventTracker.dump.writeln.fn | |
| 237 | ) | |
| 238 | | |
| 239 | def writeln(f): | |
| >>> 240 | def fn(x): | |
| 241 | return f.write(x + "\n") | |
| 242 | | |
| 243 | return fn | |
| Error (MISSING_TYPE_LINTER) Missing parameter type for torch.fx.passes.splitter_base.NodeEventTracker.dump.writeln.fn(x=) | |
| 237 | ) | |
| 238 | | |
| 239 | def writeln(f): | |
| >>> 240 | def fn(x): | |
| 241 | return f.write(x + "\n") | |
| 242 | | |
| 243 | return fn | |
| Error (MISSING_TYPE_LINTER) Missing return type for torch.fx.passes.splitter_base.NodeEventTracker.dump.dump_selected_nodes | |
| 248 | with open(self.dump_prefix + ALL_SUFFIX, "w") as f: | |
| 249 | self.print_all(writeln(f)) | |
| 250 | | |
| >>> 251 | def dump_selected_nodes(nodes): | |
| 252 | with open(self.dump_prefix + NODES_SUFFIX, "w") as f: | |
| 253 | for node_name in nodes: | |
| 254 | writeln(f"===== Tracking node {node_name} =====") | |
| Error (MISSING_TYPE_LINTER) Missing parameter type for torch.fx.passes.splitter_base.NodeEventTracker.dump.dump_selected_nodes(nodes=) | |
| 248 | with open(self.dump_prefix + ALL_SUFFIX, "w") as f: | |
| 249 | self.print_all(writeln(f)) | |
| 250 | | |
| >>> 251 | def dump_selected_nodes(nodes): | |
| 252 | with open(self.dump_prefix + NODES_SUFFIX, "w") as f: | |
| 253 | for node_name in nodes: | |
| 254 | writeln(f"===== Tracking node {node_name} =====") | |
| Successfully applied all patches. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment