Skip to content

Commit 2838e37

Browse files
authored
Minor fix for getting function's graph (#2504)
Minor fix for getting function's graph --------- Signed-off-by: Ganesan Ramalingam <[email protected]>
1 parent dc55201 commit 2838e37

File tree

1 file changed

+6
-3
lines changed

1 file changed

+6
-3
lines changed

onnxscript/rewriter/_matcher.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -183,14 +183,14 @@ def _match_value(
183183
self, pattern_value: _pattern_ir.ValuePattern, value: ir.Value | None
184184
) -> bool:
185185
"""Match an IR value against a ValuePattern instance."""
186-
if value is not None and value.graph is not self._graph_or_function:
186+
if value is not None and value.graph is not self._graph:
187187
if not isinstance(
188188
pattern_value, (_pattern_ir.Var, _pattern_ir.Constant, _pattern_ir.AnyValue)
189189
):
190190
# If the pattern value is a Var, Constant, or AnyValue, we allow it to match
191191
# values from other graphs. Otherwise, we fail the match.
192192
return self.fail(
193-
f"Value {value.name} is not in the graph {self._graph_or_function.name}. "
193+
f"Value {value.name} is not in the graph {self._graph.name}. "
194194
f"Pattern matches crossing graph boundaries are not supported."
195195
)
196196
if isinstance(pattern_value, _pattern_ir.AnyValue):
@@ -362,7 +362,10 @@ def match(
362362
complications which require careful consideration.
363363
"""
364364
self._tracer = tracer
365-
self._graph_or_function = graph_or_function[0].graph
365+
if isinstance(graph_or_function, ir.Graph):
366+
self._graph: ir.Graph = graph_or_function
367+
else:
368+
self._graph = graph_or_function.graph
366369
if self.pattern.has_single_output_node:
367370
self._init_match(verbose)
368371
return self._match_single_output_node(

0 commit comments

Comments
 (0)