add graph replace

Author: ferrine

pytensor/graph/utils.py

@@ -3,7 +3,17 @@
 import traceback
 from abc import ABCMeta
 from io import StringIO
-from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Tuple, TypeVar, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    TypeVar,
+    Union,
+)
 
 
 if TYPE_CHECKING:
@@ -420,3 +430,49 @@ def toposort(prereqs_d):
             "some orderings contain invalid elements."
         )
     return seq
+
+
+def graph_replace(
+    outputs: Sequence["Variable"], replace: Dict["Variable", "Variable"]
+) -> List["Variable"]:
+    from pytensor.graph.basic import condition_subset
+    from pytensor.graph.fg import FunctionGraph
+
+    # collect minimum graph inputs which is required to compute outputs
+    # and depend on replacements
+    conditions = list(condition_subset(outputs, replace))
+    # for the function graph we need the clean graph where
+    # inputs do not have owners
+    # this is exactly the reason to clone conditions
+    equiv = {c: c.clone() for c in conditions}
+    # some replace keys may dissapear
+    # the reason is they are inside the graph
+    # and depend on some vars in conditions
+    # but we need to keep references to make replacements
+    # we clone the replace keys to get them
+    equiv.update({r: r.clone() for r in replace})
+    # clone the graph but preserve the equiv mapping
+    fg = FunctionGraph(
+        conditions,
+        outputs,
+        # clone_get_equiv kwargs
+        copy_orphans=False,
+        copy_inputs=False,
+        memo=equiv,
+    )
+    # replace the conditions back
+    fg_replace = {equiv[c]: c for c in conditions}
+    # add the replacements on top of input mappings
+    fg_replace.update({equiv[r]: v for r, v in replace.items()})
+    # replacements have to be done in reverse topological order so that nested
+    # expressions get recursively replaced correctly
+    toposort = fg.toposort()
+    sorted_replacements = sorted(
+        tuple(fg_replace.items()),
+        # sort based on the fg toposort, if a variable has no owner, it goes first
+        key=lambda pair: (toposort.index(pair[0].owner) if pair[0].owner else -1),
+        reverse=True,
+    )
+    fg.replace_all(sorted_replacements, import_missing=True)
+
+    return list(fg.outputs)

tests/graph/test_utils.py

@@ -1,5 +1,7 @@
 import pytensor
+from pytensor.graph.utils import graph_replace
 from pytensor.tensor.type import vector
+from tests.graph.utils import MyOp, MyVariable
 
 
 def test_stack_trace():
@@ -12,3 +14,24 @@ def test_stack_trace():
         v = vector()
         assert len(v.tag.trace) == 1
         assert len(v.tag.trace[0]) == 2
+
+
+def test_replacements():
+    x = MyVariable("x")
+    y = MyVariable("y")
+    x2 = MyOp("xop")(x)
+    x2.name = "x2"
+    y2 = MyOp("yop")(y)
+    y2.name = "y2"
+
+    yc = graph_replace([x2], {x: y2})[0]
+    assert yc.owner.inputs[0] is y2
+
+    # the case where inputs have to be replaed in reverse topological order
+    o = MyOp("xyop")(x2, y2)
+    new_x = x.clone()
+    new_y2 = y2.clone()
+
+    oc = graph_replace([o], {x: new_x, y2: new_y2})[0]
+    assert oc.owner.inputs[1] is new_y2
+    assert oc.owner.inputs[0].owner.inputs[0] is new_x