Skip to content

Commit 4c4f7a0

Browse files
authored
Use logging instead of print in RewritePass (#2835)
Makes it possible to filter this out using the regular logging mechanics in Python.
1 parent 65a59e3 commit 4c4f7a0

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

onnxscript/rewriter/__init__.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
# Licensed under the MIT License.
33
from __future__ import annotations
44

5+
import logging
56
from typing import Sequence, TypeVar, Union
67

78
__all__ = [
@@ -48,6 +49,8 @@
4849
_remove_optional_bias,
4950
)
5051

52+
logger = logging.getLogger(__name__)
53+
5154
_ModelProtoOrIr = TypeVar("_ModelProtoOrIr", onnx.ModelProto, ir.Model)
5255
_DEFAULT_REWRITE_RULES: tuple[pattern.RewriteRule, ...] = (
5356
*_no_op.rules, # TODO: merge this rule into constant folding?
@@ -82,7 +85,8 @@ def __init__(
8285
def call(self, model: ir.Model) -> ir.passes.PassResult:
8386
count = self.rules.apply_to_model(model)
8487
if count:
85-
print(f"Applied {count} of general pattern rewrite rules.")
88+
logger.info("Applied %s of general pattern rewrite rules.", count)
89+
8690
return ir.passes.PassResult(model, bool(count))
8791

8892

0 commit comments

Comments
 (0)