Make the intermediate output capturer more extensible

Differential Revision: D75699297

Pull Request resolved: #11289

Author: Juntian777

devtools/inspector/_intermediate_output_capturer.py

@@ -7,24 +7,57 @@
 # pyre-unsafe
 
 
-from typing import Any, Dict, Tuple
+from typing import Any, Dict, List, Tuple
 
 import torch
 from torch.fx import GraphModule
 from torch.fx.interpreter import Interpreter
 
 
+class NodeFilter:
+    """
+    A class used to filter nodes based on extensible criteria.
+    Attributes:
+        metadata_key (str): The key to look for in the node's metadata.
+        op_type (str): The operation code to match.
+        exclude_ops (List[str]): A list of operations to exclude from the filter.
+    """
+
+    def __init__(self, metadata_key: str, op_type: str, exclude_ops: List[str] = None):
+        self.metadata_key = metadata_key
+        self.op_type = op_type
+        self.exclude_ops = exclude_ops
+
+    def matches(self, node: torch.fx.Node) -> bool:
+        return (
+            node.meta.get(self.metadata_key) is not None
+            and node.op == self.op_type
+            and all(exclude_name not in node.name for exclude_name in self.exclude_ops)
+        )
+
+
 class IntermediateOutputCapturer(Interpreter):
+    """
+    A class that captures intermediate outputs from a PyTorch graph module.
+    Attributes:
+        module (GraphModule): The graph module to capture outputs from.
+        node_filters (List[NodeFilter]): A list of filters to apply to the nodes.
+    """
+
     def __init__(self, module: GraphModule):
         super().__init__(module)
+        self.node_filters = [
+            NodeFilter("debug_handle", "call_function", exclude_ops=["getitem"])
+        ]
 
+    # Runs the graph module and captures the intermediate outputs.
     def run_and_capture(self, *args, **kwargs) -> Dict[Tuple[int, ...], Any]:
         captured_outputs = {}
 
         def capture_run_node(n: torch.fx.Node) -> Any:
             result = super(IntermediateOutputCapturer, self).run_node(n)
-            debug_handle = n.meta.get("debug_handle", None)
-            if debug_handle is not None and n.op == "call_function":
+            if all(filter.matches(n) for filter in self.node_filters):
+                debug_handle = n.meta["debug_handle"]
                 # Convert the debug handle to a tuple to use as a dictionary key
                 key = (
                     (debug_handle,)

devtools/inspector/tests/intermediate_output_capturer_test.py

@@ -111,8 +111,6 @@ def test_capture_correct_outputs(self):
             (19,): torch.tensor([[3.6000, 4.5067]]),
             (20,): torch.tensor([[0.9734, 0.9891]]),
             (21,): [torch.tensor([[0.9734]]), torch.tensor([[0.9891]])],
-            (22,): torch.tensor([[0.9734]]),
-            (23,): torch.tensor([[0.9891]]),
         }
         self.assertEqual(
             len(self.intermediate_outputs), len(expected_outputs_with_handles)