diff --git a/pytensor/configdefaults.py b/pytensor/configdefaults.py index f8354a2301..8ad3fa61f4 100644 --- a/pytensor/configdefaults.py +++ b/pytensor/configdefaults.py @@ -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", ( diff --git a/pytensor/configparser.py b/pytensor/configparser.py index 4f71e85240..328f593fe8 100644 --- a/pytensor/configparser.py +++ b/pytensor/configparser.py @@ -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 diff --git a/pytensor/graph/features.py b/pytensor/graph/features.py index 06be6d013a..3521d3b1ba 100644 --- a/pytensor/graph/features.py +++ b/pytensor/graph/features.py @@ -567,6 +567,13 @@ def replace_all_validate( 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 @@ def replace_all_validate( ) 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}" ) diff --git a/pytensor/graph/fg.py b/pytensor/graph/fg.py index e9b676f51a..bdaefc42f8 100644 --- a/pytensor/graph/fg.py +++ b/pytensor/graph/fg.py @@ -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) diff --git a/pytensor/graph/rewriting/basic.py b/pytensor/graph/rewriting/basic.py index b91e743bb6..0719870205 100644 --- a/pytensor/graph/rewriting/basic.py +++ b/pytensor/graph/rewriting/basic.py @@ -1305,9 +1305,15 @@ def transform(self, fgraph, node): 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( diff --git a/tests/graph/test_fg.py b/tests/graph/test_fg.py index e82a59e790..54ec654095 100644 --- a/tests/graph/test_fg.py +++ b/tests/graph/test_fg.py @@ -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