Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion pytensor/configdefaults.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
(
Expand Down
1 change: 1 addition & 0 deletions pytensor/configparser.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
9 changes: 8 additions & 1 deletion pytensor/graph/features.py
Original file line number Diff line number Diff line change
Expand Up @@ -567,6 +567,13 @@
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(

Check warning on line 573 in pytensor/graph/features.py

View check run for this annotation

Codecov / codecov/patch

pytensor/graph/features.py#L573

Added line #L573 was not covered by tests
","
)

for r, new_r in replacements:
try:
fgraph.replace(r, new_r, reason=reason, verbose=False, **kwargs)
Expand Down Expand Up @@ -608,7 +615,7 @@
)
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}"
)
Expand Down
14 changes: 11 additions & 3 deletions pytensor/graph/fg.py
Original file line number Diff line number Diff line change
Expand Up @@ -490,10 +490,18 @@ def replace(
"""
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)

Expand Down
12 changes: 9 additions & 3 deletions pytensor/graph/rewriting/basic.py
Original file line number Diff line number Diff line change
Expand Up @@ -1305,9 +1305,15 @@
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(

Check warning on line 1310 in pytensor/graph/rewriting/basic.py

View check run for this annotation

Codecov / codecov/patch

pytensor/graph/rewriting/basic.py#L1310

Added line #L1310 was not covered by tests
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(
Expand Down
40 changes: 40 additions & 0 deletions tests/graph/test_fg.py
Original file line number Diff line number Diff line change
Expand Up @@ -731,3 +731,43 @@ def test_dprint(self):
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