Align model dim_lengths with cloned variables in model_from_fgraph

Author: ricardoV94

pymc/model/fgraph.py

@@ -296,11 +296,20 @@ def first_non_model_var(var):
     model = Model()
     if model.parent is not None:
         raise RuntimeError("model_to_fgraph cannot be called inside a PyMC model context")
-    model._coords = getattr(fgraph, "_coords", {})
-    model._dim_lengths = getattr(fgraph, "_dim_lengths", {})
+
+    _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()}
+
+    model._coords = _coords
+    model._dim_lengths = _dim_lengths
 
     # Replace dummy `ModelVar` Ops by the underlying variables,
-    fgraph = fgraph.clone()
     model_dummy_vars = [
         model_node.outputs[0]
         for model_node in fgraph.toposort()

tests/model/test_fgraph.py

@@ -16,7 +16,7 @@
 import pytest
 
 from pytensor import config, shared
-from pytensor.graph import Constant, FunctionGraph, node_rewriter
+from pytensor.graph import Constant, FunctionGraph, graph_inputs, node_rewriter
 from pytensor.graph.rewriting.basic import in2out
 from pytensor.tensor.exceptions import NotScalarConstantError
 
@@ -164,6 +164,11 @@ def test_data(inline_views):
     np.testing.assert_allclose(pm.draw(m_new["mu"]), [100.0, 200.0])
     np.testing.assert_allclose(pm.draw(m_old["mu"]), [0.0, 1.0, 2.0], atol=1e-6)
 
+    # Check model dim_lengths contains the exact variables used in the graph of RVs
+    m_new_size_param = m_new["obs"].owner.inputs[1]
+    [m_new_dim_len] = graph_inputs([m_new_size_param])
+    assert m_new.dim_lengths["test_dim"] is m_new_dim_len
+
 
 @config.change_flags(floatX="float64")  # Avoid downcasting Ops in the graph
 def test_shared_variable():

tests/model/transform/test_optimization.py

@@ -17,7 +17,7 @@
 from pytensor.compile import SharedVariable
 from pytensor.graph import Constant
 
-from pymc import Deterministic
+from pymc import Deterministic, do
 from pymc.data import Data
 from pymc.distributions import HalfNormal, Normal
 from pymc.model import Model
@@ -132,3 +132,15 @@ def test_freeze_dims_and_data_subset():
     assert isinstance(new_m.dim_lengths["dim2"], SharedVariable)
     assert isinstance(new_m["data1"], SharedVariable)
     assert isinstance(new_m["data2"], Constant) and np.all(new_m["data2"].data == [1, 2, 3, 4, 5])
+
+
+def test_freeze_dim_after_do_intervention():
+    with Model(coords={"test_dim": range(5)}) as m:
+        mu = Data("mu", [0, 1, 2, 3, 4], dims="test_dim")
+        x = Normal("x", mu=mu, dims="test_dim")
+
+    do_m = do(m, {mu: mu * 100})
+    assert do_m["x"].type.shape == (None,)
+
+    frozen_do_m = freeze_dims_and_data(do_m)
+    assert frozen_do_m["x"].type.shape == (5,)