Allow inlining of Deterministics and Data in fgraph IR

Author: ricardoV94

pymc_experimental/tests/utils/test_model_fgraph.py

@@ -76,28 +76,37 @@ def test_basic():
     )
 
 
-def test_data():
+@pytest.mark.parametrize("inline_views", (False, True))
+def test_data(inline_views):
     """Test shared RNGs, MutableData, ConstantData and Dim lengths are handled correctly.
 
     Everything should be preserved across new and old models, except for shared RNGs
     """
     with pm.Model(coords_mutable={"test_dim": range(3)}) as m_old:
         x = pm.MutableData("x", [0.0, 1.0, 2.0], dims=("test_dim",))
         y = pm.MutableData("y", [10.0, 11.0, 12.0], dims=("test_dim",))
-        b0 = pm.ConstantData("b0", 0.0)
+        b0 = pm.ConstantData("b0", np.zeros(3))
         b1 = pm.Normal("b1")
         mu = pm.Deterministic("mu", b0 + b1 * x, dims=("test_dim",))
         obs = pm.Normal("obs", mu, sigma=1e-5, observed=y, dims=("test_dim",))
 
-    m_fgraph, memo = fgraph_from_model(m_old)
+    m_fgraph, memo = fgraph_from_model(m_old, inlined_views=inline_views)
     assert isinstance(memo[x].owner.op, ModelNamed)
     assert isinstance(memo[y].owner.op, ModelNamed)
     assert isinstance(memo[b0].owner.op, ModelNamed)
+    mu_val = memo[mu].owner.inputs[0]
+    if not inline_views:
+        # Add(b0, Mul(FreeRV(b1), x) not Add(Named(b0), Mul(FreeRV(b1), Named(x))
+        assert mu_val.owner.inputs[0] is memo[b0].owner.inputs[0]
+        assert mu_val.owner.inputs[1].owner.inputs[1] is memo[x].owner.inputs[0]
+    else:
+        assert mu_val.owner.inputs[0] is memo[b0]
+        assert mu_val.owner.inputs[1].owner.inputs[1] is memo[x]
 
     m_new = model_from_fgraph(m_fgraph)
 
     # ConstantData is preserved
-    assert m_new["b0"].data == m_old["b0"].data
+    assert np.all(m_new["b0"].data == m_old["b0"].data)
 
     # Shared non-rng shared variables are preserved
     assert m_new["x"].container is x.container
@@ -114,7 +123,8 @@ def test_data():
     np.testing.assert_array_almost_equal(pm.draw(m_new["x"]), [100.0, 200.0])
 
 
-def test_deterministics():
+@pytest.mark.parametrize("inline_views", (False, True))
+def test_deterministics(inline_views):
     """Test handling of deterministics.
 
     We don't want Deterministics in the middle of the FunctionGraph, as they would make rewrites cumbersome
@@ -140,22 +150,27 @@ 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, inlined_views=inline_views)
 
     # 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
     # [Det(mu), Det(sigma)]
     mu = det_mu.owner.inputs[0]
     sigma = det_sigma.owner.inputs[0]
-    # [FreeRV(y(mu, sigma))] not [FreeRV(y(Det(mu), Det(sigma)))]
-    assert y.owner.inputs[0].owner.inputs[3] is mu
     assert y.owner.inputs[0].owner.inputs[4] is sigma
-    # [FreeRV(z(y))] not [FreeRV(z(Det(Det(y))))]
-    assert z.owner.inputs[0].owner.inputs[3] is y
-    # [Det(y), Det(y)], not [Det(y), Det(Det(y))]
-    assert det_y_.owner.inputs[0] is y
-    assert det_y__.owner.inputs[0] is y
     assert det_y_ is not det_y__
+    assert det_y_.owner.inputs[0] is y
+    if not inline_views:
+        # FreeRV(y(mu, sigma)) not FreeRV(y(Det(mu), Det(sigma)))
+        assert y.owner.inputs[0].owner.inputs[3] is mu
+        # FreeRV(z(y)) not FreeRV(z(Det(Det(y))))
+        assert z.owner.inputs[0].owner.inputs[3] is y
+        # Det(y), not Det(Det(y))
+        assert det_y__.owner.inputs[0] is y
+    else:
+        assert y.owner.inputs[0].owner.inputs[3] is det_mu
+        assert z.owner.inputs[0].owner.inputs[3] is det_y__
+        assert det_y__.owner.inputs[0] is det_y_
 
     # Both mu and sigma deterministics are now in the graph of x to y
     m = model_from_fgraph(fg)

pymc_experimental/utils/model_fgraph.py

@@ -109,11 +109,20 @@ def local_remove_identity(fgraph, node):
 remove_identity_rewrite = out2in(local_remove_identity)
 
 
-def fgraph_from_model(model: Model) -> Tuple[FunctionGraph, Dict[Variable, Variable]]:
+def fgraph_from_model(
+    model: Model, inlined_views=False
+) -> Tuple[FunctionGraph, Dict[Variable, Variable]]:
     """Convert Model to FunctionGraph.
 
     See: model_from_fgraph
 
+    Parameters
+    ----------
+    model: PyMC model
+    inlined_views: bool, default False
+        Whether "view" variables (Deterministics and Data) should be inlined among RVs in the fgraph,
+        or show up as separate branches.
+
     Returns
     -------
     fgraph: FunctionGraph
@@ -141,13 +150,15 @@ def fgraph_from_model(model: Model) -> Tuple[FunctionGraph, Dict[Variable, Varia
     # We copy Deterministics (Identity Op) so that they don't show in between "main" variables
     # We later remove these Identity Ops when we have a Deterministic ModelVar Op as a separator
     old_deterministics = model.deterministics
-    deterministics = [det.copy(det.name) for det in old_deterministics]
+    deterministics = [det if inlined_views else det.copy(det.name) for det in old_deterministics]
     # Other variables that are in model.named_vars but are not any of the categories above
     # E.g., MutableData, ConstantData, _dim_lengths
     # We use the same trick as deterministics!
     accounted_for = free_rvs + observed_rvs + potentials + old_deterministics
     old_other_named_vars = [var for var in model.named_vars.values() if var not in accounted_for]
-    other_named_vars = [var.copy(var.name) for var in old_other_named_vars]
+    other_named_vars = [
+        var if inlined_views else var.copy(var.name) for var in old_other_named_vars
+    ]
     value_vars = [val for val in rvs_to_values.values() if val not in old_other_named_vars]
 
     model_vars = rvs + potentials + deterministics + other_named_vars + value_vars
@@ -211,7 +222,7 @@ def fgraph_from_model(model: Model) -> Tuple[FunctionGraph, Dict[Variable, Varia
     # Reference model vars in memo
     inverse_memo = {v: k for k, v in memo.items()}
     for var, model_var in replacements:
-        if isinstance(
+        if not inlined_views and isinstance(
             model_var.owner is not None and model_var.owner.op, (ModelDeterministic, ModelNamed)
         ):
             # Ignore extra identity that will be removed at the end