Return memo dictionary in fgraph_from_model

Author: ricardoV94

pymc_experimental/tests/utils/test_model_fgraph.py

@@ -7,7 +7,11 @@
 from pytensor.tensor.exceptions import NotScalarConstantError
 
 from pymc_experimental.utils.model_fgraph import (
+    ModelDeterministic,
     ModelFreeRV,
+    ModelNamed,
+    ModelObservedRV,
+    ModelPotential,
     ModelVar,
     fgraph_from_model,
     model_deterministic,
@@ -23,11 +27,17 @@ def test_basic():
         y = pm.Deterministic("y", x + 1)
         w = pm.HalfNormal("w", pm.math.exp(y))
         z = pm.Normal("z", y, w, observed=[0, 1, 2], dims=("test_dim",))
-        pm.Potential("pot", x * 2)
+        pot = pm.Potential("pot", x * 2)
 
-    m_fgraph = fgraph_from_model(m_old)
+    m_fgraph, memo = fgraph_from_model(m_old)
     assert isinstance(m_fgraph, FunctionGraph)
 
+    assert isinstance(memo[x].owner.op, ModelFreeRV)
+    assert isinstance(memo[y].owner.op, ModelDeterministic)
+    assert isinstance(memo[w].owner.op, ModelFreeRV)
+    assert isinstance(memo[z].owner.op, ModelObservedRV)
+    assert isinstance(memo[pot].owner.op, ModelPotential)
+
     m_new = model_from_fgraph(m_fgraph)
     assert isinstance(m_new, pm.Model)
 
@@ -79,7 +89,12 @@ def test_data():
         mu = pm.Deterministic("mu", b0 + b1 * x, dims=("test_dim",))
         obs = pm.Normal("obs", mu, sigma=1e-5, observed=y, dims=("test_dim",))
 
-    m_new = model_from_fgraph(fgraph_from_model(m_old))
+    m_fgraph, memo = fgraph_from_model(m_old)
+    assert isinstance(memo[x].owner.op, ModelNamed)
+    assert isinstance(memo[y].owner.op, ModelNamed)
+    assert isinstance(memo[b0].owner.op, ModelNamed)
+
+    m_new = model_from_fgraph(m_fgraph)
 
     # ConstantData is preserved
     assert m_new["b0"].data == m_old["b0"].data
@@ -125,7 +140,7 @@ def test_deterministics():
     assert m["y"].owner.inputs[3] is m["mu"]
     assert m["y"].owner.inputs[4] is not m["sigma"]
 
-    fg = fgraph_from_model(m)
+    fg, _ = fgraph_from_model(m)
 
     # Check that no Deterministics are in graph of x to y and y to z
     x, y, z, det_mu, det_sigma, det_y_, det_y__ = fg.outputs
@@ -173,7 +188,7 @@ def test_sub_model_error():
         with pm.Model() as sub_m:
             y = pm.Normal("y", x)
 
-    nodes = [v for v in fgraph_from_model(m).toposort() if not isinstance(v.op, ModelVar)]
+    nodes = [v for v in fgraph_from_model(m)[0].toposort() if not isinstance(v.op, ModelVar)]
     assert len(nodes) == 2
     assert isinstance(nodes[0].op, pm.Beta)
     assert isinstance(nodes[1].op, pm.Normal)
@@ -234,7 +249,7 @@ def test_fgraph_rewrite(non_centered_rewrite):
         subject_mean = pm.Normal("subject_mean", group_mean, group_std, dims=("subject",))
         obs = pm.Normal("obs", subject_mean, 1, observed=np.zeros(10), dims=("subject",))
 
-    fg = fgraph_from_model(m_old)
+    fg, _ = fgraph_from_model(m_old)
     non_centered_rewrite.apply(fg)
 
     m_new = model_from_fgraph(fg)

pymc_experimental/utils/model_fgraph.py

@@ -1,4 +1,4 @@
-from typing import Optional, Sequence, Tuple
+from typing import Dict, Optional, Sequence, Tuple
 
 import pytensor
 from pymc.logprob.transforms import RVTransform
@@ -109,16 +109,19 @@ def local_remove_identity(fgraph, node):
 remove_identity_rewrite = out2in(local_remove_identity)
 
 
-def fgraph_from_model(model: Model) -> FunctionGraph:
+def fgraph_from_model(model: Model) -> Tuple[FunctionGraph, Dict[Variable, Variable]]:
     """Convert Model to FunctionGraph.
 
-    Create a FunctionGraph that includes a copy of model variables, wrapped in dummy `ModelVar` Ops.
-
-    PyTensor rewrites can be used to transform the FunctionGraph.
+    See: model_from_fgraph
 
-    It should be possible to reconstruct a valid PyMC model using `model_from_fgraph`.
+    Returns
+    -------
+    fgraph: FunctionGraph
+        FunctionGraph that includes a copy of model variables, wrapped in dummy `ModelVar` Ops.
+        It should be possible to reconstruct a valid PyMC model using `model_from_fgraph`.
 
-    See: model_from_fgraph
+    memo: Dict
+        A dictionary mapping original model variables to the equivalent nodes in the fgraph.
     """
 
     if any(v is not None for v in model.rvs_to_initial_values.values()):
@@ -202,7 +205,19 @@ def fgraph_from_model(model: Model) -> FunctionGraph:
             new_var = var
         new_vars.append(new_var)
 
-    toposort_replace(fgraph, tuple(zip(vars, new_vars)))
+    replacements = tuple(zip(vars, new_vars))
+    toposort_replace(fgraph, replacements)
+
+    # Reference model vars in memo
+    inverse_memo = {v: k for k, v in memo.items()}
+    for var, model_var in replacements:
+        if isinstance(
+            model_var.owner is not None and model_var.owner.op, (ModelDeterministic, ModelNamed)
+        ):
+            # Ignore extra identity that will be removed at the end
+            var = var.owner.inputs[0]
+        original_var = inverse_memo[var]
+        memo[original_var] = model_var
 
     # Remove value variable as outputs, now that they are graph inputs
     first_value_idx = len(fgraph.outputs) - len(value_vars)
@@ -212,7 +227,7 @@ def fgraph_from_model(model: Model) -> FunctionGraph:
     # Now that we have Deterministic dummy Ops, we remove the noisy `Identity`s from the graph
     remove_identity_rewrite.apply(fgraph)
 
-    return fgraph
+    return fgraph, memo
 
 
 def model_from_fgraph(fgraph: FunctionGraph) -> Model:
@@ -282,7 +297,7 @@ def model_from_fgraph(fgraph: FunctionGraph) -> Model:
     return model
 
 
-def clone_model(model: Model) -> Model:
+def clone_model(model: Model) -> Tuple[Model]:
     """Clone a PyMC model.
 
     Recreates a PyMC model with clones of the original variables.
@@ -310,4 +325,4 @@ def clone_model(model: Model) -> Model:
             z = pm.Deterministic("z", clone_x + 1)
 
     """
-    return model_from_fgraph(fgraph_from_model(model))
+    return model_from_fgraph(fgraph_from_model(model)[0])