Revert "[fx][subgraph_rewriter] Change match_filter to be a List in replace_pattern_with_filters (#87257)"

pytorchmergebot · pytorchmergebot · commit d13b6781d8b7 · 2022-10-28T17:55:19.000Z
This reverts commit 5865083. Reverted #87257 on behalf of https://github.com/weiwangmeta due to breaking internal builds/BC-breaking change
diff --git a/test/fx/test_subgraph_rewriter.py b/test/fx/test_subgraph_rewriter.py
@@ -773,7 +773,7 @@ def gemm_bias_mul_replacement_with_c(a, b, bias, c):
 
         self.assertEqual(repalcement_node_found, 2)
 
-    def test_replace_pattern_with_filters(self):
+    def test_replace_pattern_with_filter(self):
         class M(torch.nn.Module):
             def __init__(self):
                 super().__init__()
@@ -833,10 +833,10 @@ def num_repalcement_node_found(traced):
 
         # match with filter, should find 1 match
         traced = symbolic_trace(M())
-        matches = subgraph_rewriter.replace_pattern_with_filters(
+        matches = subgraph_rewriter.replace_pattern_with_filter(
             traced,
             BinaryOpScalarReLUPattern,
             BinaryOpScalarReLUReplacement,
-            [second_input_is_scalar])
+            second_input_is_scalar)
         self.assertEqual(len(matches), 1)
         self.assertEqual(num_repalcement_node_found(traced), 1)
diff --git a/torch/fx/subgraph_rewriter.py b/torch/fx/subgraph_rewriter.py
@@ -8,7 +8,7 @@
 from typing import Callable, Dict, List, NamedTuple, Optional, Set
 import torch
 
-__all__ = ['Match', 'replace_pattern', 'replace_pattern_with_filters']
+__all__ = ['Match', 'replace_pattern', 'replace_pattern_with_filter']
 
 @compatibility(is_backward_compatible=True)
 class Match(NamedTuple):
@@ -185,11 +185,11 @@ def forward(self, x, w1, w2):
 
 # Experimental API, not backward compatible
 @compatibility(is_backward_compatible=False)
-def replace_pattern_with_filters(
+def replace_pattern_with_filter(
     gm: GraphModule,
     pattern: Callable,
     replacement: Callable,
-    match_filters: List[Callable[["InternalMatch", Graph, Graph], bool]],  # type: ignore[name-defined]
+    match_filter: Callable[["InternalMatch", Graph, Graph], bool],  # type: ignore[name-defined]
 ) -> List[Match]:
     """
     See replace_pattern for documentation. This function is an overload with an additional match_filter argument.
@@ -200,21 +200,18 @@ def replace_pattern_with_filters(
             definition of InternalMatch.
     """
 
-    return _replace_pattern(gm, pattern, replacement, match_filters)
+    return _replace_pattern(gm, pattern, replacement, match_filter)
 
 
 def _replace_pattern(
     gm: GraphModule,
     pattern: Callable,
     replacement: Callable,
-    match_filters: List[Callable[["InternalMatch", Graph, Graph], bool]] = None  # type: ignore[name-defined]
+    match_filter: Optional[Callable[["InternalMatch", Graph, Graph], bool]] = None  # type: ignore[name-defined]
 ) -> List[Match]:
 
     from torch.fx.passes.utils.matcher_utils import SubgraphMatcher, InternalMatch
 
-    if match_filters is None:
-        match_filters = []
-
     # Get the graphs for `gm`, `pattern`, `replacement`
     original_graph: Graph = gm.graph
     pattern_graph: Graph = symbolic_trace(pattern).graph
@@ -225,11 +222,8 @@ def _replace_pattern(
     _matches: List[InternalMatch] = matcher.match(original_graph)
 
     # Filter out matches that don't match the filter
-    _matches = [
-        m for m in _matches
-        if all(match_filter(m, original_graph, pattern_graph)
-               for match_filter in match_filters)
-    ]
+    if match_filter:
+        _matches = [m for m in _matches if match_filter(m, original_graph, pattern_graph)]
 
     replacement_placeholders = [n for n in replacement_graph.nodes if n.op == "placeholder"]
 
