Allow ignoring specific rewrites from optimizer_verbose

Author: ricardoV94

pytensor/configdefaults.py

@@ -423,6 +423,13 @@ def add_compile_configvars():
  in_c_key=False,
  )
 
+ config.add(
+ "optimizer_verbose_ignore",
+ ("Ignore rewrites from optimizers with these names. Separate names with ','"),
+ StrParam(""),
+ in_c_key=False,
+ )
+
  config.add(
  "on_opt_error",
  (

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