Add option to avoid cloning fgraph in model_from_fgraph

In most function transforms the caller is both creating the fgraph representation and discarding it, so it's safe to mutate the fgraph in place.

Author: ricardoV94

pymc/model/fgraph.py

@@ -278,12 +278,18 @@ def fgraph_from_model(
     return fgraph, memo
 
 
-def model_from_fgraph(fgraph: FunctionGraph) -> Model:
+def model_from_fgraph(fgraph: FunctionGraph, mutate_fgraph: bool = False) -> Model:
     """Convert FunctionGraph to PyMC model.
 
-    This requires nodes to be properly tagged with `ModelVar` dummy Ops.
+    Parameters
+    ----------
+    fgraph: FunctionGraph
+        fgraph representation of a PyMC model, with dummy `ModelVar` Ops.
+        See `fgraph_from_model` for more details.
 
-    See: fgraph_from_model
+    mutate_fgraph: bool, default False
+        Whether the function is allowed to modify the fgraph (and it's variables) in place.
+         This is useful if these are not needed anymore after the model is created.
     """
 
     def first_non_model_var(var):
@@ -300,11 +306,12 @@ def first_non_model_var(var):
     _coords = getattr(fgraph, "_coords", {})
     _dim_lengths = getattr(fgraph, "_dim_lengths", {})
 
-    fgraph, memo = fgraph.clone_get_equiv(check_integrity=False, attach_feature=False)
-    # Shared dim lengths are not extracted from the fgraph representation,
-    # so we need to update after we clone the fgraph
-    # TODO: Consider representing/extracting them from the fgraph!
-    _dim_lengths = {k: memo.get(v, v) for k, v in _dim_lengths.items()}
+    if not mutate_fgraph:
+        fgraph, memo = fgraph.clone_get_equiv(check_integrity=False, attach_feature=False)
+        # Shared dim lengths are not extracted from the fgraph representation,
+        # so we need to update after we clone the fgraph
+        # TODO: Consider representing/extracting them from the fgraph!
+        _dim_lengths = {k: memo.get(v, v) for k, v in _dim_lengths.items()}
 
     model._coords = _coords
     model._dim_lengths = _dim_lengths
@@ -385,7 +392,7 @@ def clone_model(model: Model) -> Model:
             z = pm.Deterministic("z", clone_x + 1)
 
     """
-    return model_from_fgraph(fgraph_from_model(model)[0])
+    return model_from_fgraph(fgraph_from_model(model)[0], mutate_fgraph=True)
 
 
 def extract_dims(var) -> tuple:

pymc/model/transform/basic.py

@@ -50,7 +50,7 @@ def prune_vars_detached_from_observed(model: Model) -> Model:
     }
     for node_to_remove in nodes_to_remove:
         fgraph.remove_node(node_to_remove)
-    return model_from_fgraph(fgraph)
+    return model_from_fgraph(fgraph, mutate_fgraph=True)
 
 
 def parse_vars(model: Model, vars: ModelVariable | Sequence[ModelVariable]) -> list[Variable]:

pymc/model/transform/conditioning.py

@@ -117,7 +117,7 @@ def observe(
 
     toposort_replace(fgraph, tuple(replacements.items()))
 
-    return model_from_fgraph(fgraph)
+    return model_from_fgraph(fgraph, mutate_fgraph=True)
 
 
 def do(
@@ -215,7 +215,7 @@ def do(
     # Replace variables by interventions
     toposort_replace(fgraph, tuple(replacements.items()))
 
-    model = model_from_fgraph(fgraph)
+    model = model_from_fgraph(fgraph, mutate_fgraph=True)
     if prune_vars:
         return prune_vars_detached_from_observed(model)
     return model
@@ -302,7 +302,7 @@ def change_value_transforms(
         replacements[dummy_rv] = new_dummy_rv
 
     toposort_replace(fgraph, tuple(replacements.items()))
-    return model_from_fgraph(fgraph)
+    return model_from_fgraph(fgraph, mutate_fgraph=True)
 
 
 def remove_value_transforms(

pymc/model/transform/optimization.py

@@ -104,7 +104,7 @@ def freeze_dims_and_data(
         replacements[old_value] = new_value
     fg.replace_all(tuple(replacements.items()), import_missing=True)
 
-    return model_from_fgraph(fg)
+    return model_from_fgraph(fg, mutate_fgraph=True)
 
 
 __all__ = ("freeze_dims_and_data",)