Return memo dictionary in fgraph_from_model

Author: ricardoV94

pymc_experimental/tests/utils/test_model_fgraph.py

@@ -23,11 +23,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 memo[x] in m_fgraph.variables
+    assert memo[y] in m_fgraph.variables
+    assert memo[w] in m_fgraph.variables
+    assert memo[z] in m_fgraph.variables
+    assert memo[pot] in m_fgraph.variables
+
     m_new = model_from_fgraph(m_fgraph)
     assert isinstance(m_new, pm.Model)
 
@@ -79,7 +85,7 @@ 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_new = model_from_fgraph(fgraph_from_model(m_old)[0])
 
     # ConstantData is preserved
     assert m_new["b0"].data == m_old["b0"].data
@@ -125,7 +131,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 +179,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 +240,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
+from typing import Dict, Optional, Tuple
 
 import pytensor
 from pymc.logprob.transforms import RVTransform
@@ -107,16 +107,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()):
@@ -200,7 +203,17 @@ 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):
+            # 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)
@@ -210,7 +223,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:
@@ -280,7 +293,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.
@@ -308,4 +321,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])