Add reason strings to FunctionGraph.replace calls

brandonwillard · Dec 2, 2022 · f793027 · f793027
1 parent 2b48015
commit f793027
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/aeppl/rewriting.py b/aeppl/rewriting.py
@@ -326,6 +326,6 @@ def construct_ir_fgraph(
         new_to_old = tuple(
             (v, k) for k, v in rv_remapper.measurable_conversions.items()
         )
-        fgraph.replace_all(new_to_old)
+        fgraph.replace_all(new_to_old, reason="construct_ir_fgraph")
 
     return fgraph, rv_values, memo
diff --git a/aeppl/utils.py b/aeppl/utils.py
@@ -188,7 +188,9 @@ def expand_replace(var: TensorVariable) -> List[TensorVariable]:
             clone=False,
         )
 
-        fg.replace_all(replacements.items(), import_missing=True)
+        fg.replace_all(
+            replacements.items(), import_missing=True, reason="replace_rvs_in_graphs"
+        )
 
         graphs = list(fg.outputs)