Allow ignoring printing of specific rewrites from optimizer_verbose

Author: ricardoV94

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",
         (

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

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}"
             )

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)
 

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(

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