move clone_replace to a separate file

Author: ferrine

pytensor/__init__.py

@@ -73,7 +73,8 @@ def disable_log_handler(logger=pytensor_logger, handler=logging_default_handler)
 __api_version__ = 1
 
 # isort: off
-from pytensor.graph.basic import Variable, clone_replace
+from pytensor.graph.basic import Variable
+from pytensor.graph.replace import clone_replace
 
 # isort: on
 

pytensor/compile/builders.py

@@ -16,13 +16,13 @@
     Constant,
     NominalVariable,
     Variable,
-    clone_replace,
     graph_inputs,
     io_connection_pattern,
 )
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.null_type import NullType
 from pytensor.graph.op import HasInnerGraph, Op
+from pytensor.graph.replace import clone_replace
 from pytensor.graph.rewriting.basic import in2out, node_rewriter
 from pytensor.graph.utils import MissingInputError
 from pytensor.tensor.rewriting.shape import ShapeFeature

pytensor/graph/__init__.py

@@ -7,9 +7,9 @@
     Constant,
     graph_inputs,
     clone,
-    clone_replace,
     ancestors,
 )
+from pytensor.graph.replace import clone_replace
 from pytensor.graph.op import Op
 from pytensor.graph.type import Type
 from pytensor.graph.fg import FunctionGraph

pytensor/graph/basic.py

@@ -1151,53 +1151,6 @@ def clone_get_equiv(
     return memo
 
 
-def clone_replace(
-    output: Collection[Variable],
-    replace: Optional[
-        Union[Iterable[Tuple[Variable, Variable]], Dict[Variable, Variable]]
-    ] = None,
-    **rebuild_kwds,
-) -> List[Variable]:
-    """Clone a graph and replace subgraphs within it.
-
-    It returns a copy of the initial subgraph with the corresponding
-    substitutions.
-
-    Parameters
-    ----------
-    output
-        PyTensor expression that represents the computational graph.
-    replace
-        Dictionary describing which subgraphs should be replaced by what.
-    rebuild_kwds
-        Keywords to `rebuild_collect_shared`.
-
-    """
-    from pytensor.compile.function.pfunc import rebuild_collect_shared
-
-    items: Union[List[Tuple[Variable, Variable]], Tuple[Tuple[Variable, Variable], ...]]
-    if isinstance(replace, dict):
-        items = list(replace.items())
-    elif isinstance(replace, (list, tuple)):
-        items = replace
-    elif replace is None:
-        items = []
-    else:
-        raise ValueError(
-            "replace is neither a dictionary, list, "
-            f"tuple or None ! The value provided is {replace},"
-            f"of type {type(replace)}"
-        )
-    tmp_replace = [(x, x.type()) for x, y in items]
-    new_replace = [(x, y) for ((_, x), (_, y)) in zip(tmp_replace, items)]
-    _, _outs, _ = rebuild_collect_shared(output, [], tmp_replace, [], **rebuild_kwds)
-
-    # TODO Explain why we call it twice ?!
-    _, outs, _ = rebuild_collect_shared(_outs, [], new_replace, [], **rebuild_kwds)
-
-    return cast(List[Variable], outs)
-
-
 def general_toposort(
     outputs: Iterable[T],
     deps: Callable[[T], Union[OrderedSet, List[T]]],

pytensor/graph/replace.py

@@ -0,0 +1,60 @@
+from typing import (
+    Collection,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+    cast,
+)
+
+from pytensor.graph.basic import Constant, Variable
+
+
+def clone_replace(
+    output: Collection[Variable],
+    replace: Optional[
+        Union[Iterable[Tuple[Variable, Variable]], Dict[Variable, Variable]]
+    ] = None,
+    **rebuild_kwds,
+) -> List[Variable]:
+    """Clone a graph and replace subgraphs within it.
+
+    It returns a copy of the initial subgraph with the corresponding
+    substitutions.
+
+    Parameters
+    ----------
+    output
+        PyTensor expression that represents the computational graph.
+    replace
+        Dictionary describing which subgraphs should be replaced by what.
+    rebuild_kwds
+        Keywords to `rebuild_collect_shared`.
+
+    """
+    from pytensor.compile.function.pfunc import rebuild_collect_shared
+
+    items: Union[List[Tuple[Variable, Variable]], Tuple[Tuple[Variable, Variable], ...]]
+    if isinstance(replace, dict):
+        items = list(replace.items())
+    elif isinstance(replace, (list, tuple)):
+        items = replace
+    elif replace is None:
+        items = []
+    else:
+        raise ValueError(
+            "replace is neither a dictionary, list, "
+            f"tuple or None ! The value provided is {replace},"
+            f"of type {type(replace)}"
+        )
+    tmp_replace = [(x, x.type()) for x, y in items]
+    new_replace = [(x, y) for ((_, x), (_, y)) in zip(tmp_replace, items)]
+    _, _outs, _ = rebuild_collect_shared(output, [], tmp_replace, [], **rebuild_kwds)
+
+    # TODO Explain why we call it twice ?!
+    _, outs, _ = rebuild_collect_shared(_outs, [], new_replace, [], **rebuild_kwds)
+
+    return cast(List[Variable], outs)

pytensor/ifelse.py

@@ -20,8 +20,9 @@
 from pytensor import as_symbolic
 from pytensor.compile import optdb
 from pytensor.configdefaults import config
-from pytensor.graph.basic import Apply, Variable, clone_replace, is_in_ancestors
+from pytensor.graph.basic import Apply, Variable, is_in_ancestors
 from pytensor.graph.op import _NoPythonOp
+from pytensor.graph.replace import clone_replace
 from pytensor.graph.rewriting.basic import GraphRewriter, in2out, node_rewriter
 from pytensor.graph.type import HasDataType, HasShape
 from pytensor.tensor.shape import Reshape, Shape, SpecifyShape, Unbroadcast

pytensor/scan/basic.py

@@ -6,8 +6,9 @@
 from pytensor.compile.function.pfunc import construct_pfunc_ins_and_outs
 from pytensor.compile.sharedvalue import SharedVariable, collect_new_shareds
 from pytensor.configdefaults import config
-from pytensor.graph.basic import Constant, Variable, clone_replace, graph_inputs
+from pytensor.graph.basic import Constant, Variable, graph_inputs
 from pytensor.graph.op import get_test_value
+from pytensor.graph.replace import clone_replace
 from pytensor.graph.utils import MissingInputError, TestValueError
 from pytensor.scan.op import Scan, ScanInfo
 from pytensor.scan.utils import expand_empty, safe_new, until

pytensor/scan/op.py

@@ -65,13 +65,13 @@
 from pytensor.graph.basic import (
     Apply,
     Variable,
-    clone_replace,
     equal_computations,
     graph_inputs,
     io_connection_pattern,
 )
 from pytensor.graph.features import NoOutputFromInplace
 from pytensor.graph.op import HasInnerGraph, Op
+from pytensor.graph.replace import clone_replace
 from pytensor.graph.utils import InconsistencyError, MissingInputError
 from pytensor.link.c.basic import CLinker
 from pytensor.link.c.exceptions import MissingGXX

pytensor/scan/rewriting.py

@@ -18,7 +18,6 @@
     Apply,
     Constant,
     Variable,
-    clone_replace,
     equal_computations,
     graph_inputs,
     io_toposort,
@@ -28,6 +27,7 @@
 from pytensor.graph.features import ReplaceValidate
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import compute_test_value
+from pytensor.graph.replace import clone_replace
 from pytensor.graph.rewriting.basic import GraphRewriter, in2out, node_rewriter
 from pytensor.graph.rewriting.db import EquilibriumDB, SequenceDB
 from pytensor.graph.type import HasShape

pytensor/scan/utils.py

@@ -14,14 +14,9 @@
 from pytensor import tensor as at
 from pytensor.compile.profiling import ProfileStats
 from pytensor.configdefaults import config
-from pytensor.graph.basic import (
-    Constant,
-    Variable,
-    clone_replace,
-    equal_computations,
-    graph_inputs,
-)
+from pytensor.graph.basic import Constant, Variable, equal_computations, graph_inputs
 from pytensor.graph.op import get_test_value
+from pytensor.graph.replace import clone_replace
 from pytensor.graph.type import HasDataType
 from pytensor.graph.utils import TestValueError
 from pytensor.tensor.basic import AllocEmpty, cast