提交 7f031251 authored 作者: Ricardo Vieira's avatar Ricardo Vieira 提交者: Ricardo Vieira

Allow ignoring printing of specific rewrites from optimizer_verbose

上级 5008fab7
......@@ -418,11 +418,18 @@ def add_compile_configvars():
config.add(
"optimizer_verbose",
"If True, we print all optimization being applied",
"Print information about rewrites that are applied during a graph transformation.",
BoolParam(False),
in_c_key=False,
)
config.add(
"optimizer_verbose_ignore",
"Do not print information for rewrites with these names when `optimizer_verbose` is `True`. Separate names with ','",
StrParam(""),
in_c_key=False,
)
config.add(
"on_opt_error",
(
......
......@@ -81,6 +81,7 @@ class PyTensorConfigParser:
allow_gc: bool
optimizer: str
optimizer_verbose: bool
optimizer_verbose_ignore: str
on_opt_error: str
nocleanup: bool
on_unused_input: str
......
......@@ -567,6 +567,13 @@ class ReplaceValidate(History, Validator):
if verbose is None:
verbose = config.optimizer_verbose
if verbose:
print_reason = True
if config.optimizer_verbose_ignore:
print_reason = str(reason) not in config.optimizer_verbose_ignore.split(
","
)
for r, new_r in replacements:
try:
fgraph.replace(r, new_r, reason=reason, verbose=False, **kwargs)
......@@ -608,7 +615,7 @@ class ReplaceValidate(History, Validator):
)
raise
if verbose:
if verbose and print_reason:
print( # noqa: T201
f"rewriting: rewrite {reason} replaces {r} of {r.owner} with {new_r} of {new_r.owner}"
)
......
......@@ -490,10 +490,18 @@ class FunctionGraph(MetaObject):
"""
if verbose is None:
verbose = config.optimizer_verbose
if verbose:
print( # noqa: T201
f"rewriting: rewrite {reason} replaces {var} of {var.owner} with {new_var} of {new_var.owner}"
)
print_reason = True
if config.optimizer_verbose_ignore:
print_reason = str(reason) not in config.optimizer_verbose_ignore.split(
","
)
if print_reason:
print( # noqa: T201
f"rewriting: rewrite {reason} replaces {var} of {var.owner} with {new_var} of {new_var.owner}"
)
new_var = var.type.filter_variable(new_var, allow_convert=True)
......
......@@ -1305,9 +1305,15 @@ class SequentialNodeRewriter(NodeRewriter):
new_vars = list(new_repl.values())
if config.optimizer_verbose:
print( # noqa: T201
f"rewriting: rewrite {rewrite} replaces node {node} with {new_repl}"
)
print_reason = True
if config.optimizer_verbose_ignore:
print_reason = str(
rewrite
) not in config.optimizer_verbose_ignore.split(",")
if print_reason:
print( # noqa: T201
f"rewriting: rewrite {rewrite} replaces node {node} with {new_repl}"
)
if self.profile:
self.node_created[rewrite] += len(
......
......@@ -731,3 +731,43 @@ class TestFunctionGraph:
o1 = op1(r1, r2)
fg = FunctionGraph([r1, r2], [o1], clone=False)
assert fg.dprint(file="str") == debugprint(fg, file="str")
def test_optimizer_verbose(self, capsys):
x = MyVariable("x")
y = MyVariable("y")
z = MyVariable("z")
o1 = op1(x, y)
fgraph = FunctionGraph([x, y, z], [o1], clone=False)
with config.change_flags(optimizer_verbose=False):
fgraph.replace(y, z, reason="y->z")
cap_out = capsys.readouterr().out
assert cap_out == ""
with config.change_flags(optimizer_verbose=True):
fgraph.replace(z, y, reason="z->y")
cap_out = capsys.readouterr().out
assert "z->y" in cap_out
with config.change_flags(
optimizer_verbose=True, optimizer_verbose_ignore="y->z"
):
fgraph.replace(y, z, reason="y->z")
fgraph.replace(z, y, reason="z->y")
cap_out = capsys.readouterr().out
assert "y->z" not in cap_out
assert "z->y" in cap_out
with config.change_flags(
optimizer_verbose=True, optimizer_verbose_ignore="y->z,z->y"
):
fgraph.replace(y, z, reason="y->z")
fgraph.replace(z, y, reason="z->y")
cap_out = capsys.readouterr().out
assert "y->z" not in cap_out
assert "z->y" not in cap_out
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论