Ignore constant shape dependencies in MarginalModel

Author: ricardoV94

pymc_experimental/marginal_model.py

@@ -21,6 +21,8 @@
 
 __all__ = ["MarginalModel"]
 
+from pytensor.tensor.shape import Shape
+
 
 class MarginalModel(Model):
     """Subclass of PyMC Model that implements functionality for automatic
@@ -251,9 +253,24 @@ class FiniteDiscreteMarginalRV(MarginalRV):
     """Base class for Finite Discrete Marginalized RVs"""
 
 
+def static_shape_ancestors(vars):
+    """Identify ancestors Shape Ops of static shapes (therefore constant in a valid graph)."""
+    return [
+        var
+        for var in ancestors(vars)
+        if (
+            var.owner
+            and isinstance(var.owner.op, Shape)
+            # All static dims lengths of Shape input are known
+            and None not in var.owner.inputs[0].type.shape
+        )
+    ]
+
+
 def find_conditional_input_rvs(output_rvs, all_rvs):
-    """Find conditionally indepedent input RVs"""
+    """Find conditionally indepedent input RVs."""
     blockers = [other_rv for other_rv in all_rvs if other_rv not in output_rvs]
+    blockers += static_shape_ancestors(tuple(all_rvs) + tuple(output_rvs))
     return [
         var
         for var in ancestors(output_rvs, blockers=blockers)

pymc_experimental/tests/test_marginal_model.py

@@ -12,7 +12,11 @@
 from pymc.util import UNSET
 from scipy.special import logsumexp
 
-from pymc_experimental.marginal_model import FiniteDiscreteMarginalRV, MarginalModel
+from pymc_experimental.marginal_model import (
+    FiniteDiscreteMarginalRV,
+    MarginalModel,
+    is_conditional_dependent,
+)
 
 
 @pytest.fixture
@@ -411,3 +415,14 @@ def test_marginalized_transforms(transform, expected_warning):
             transform_name = transform.name
         assert f"sigma_{transform_name}__" in ip
     np.testing.assert_allclose(m.compile_logp()(ip), m_ref.compile_logp()(ip))
+
+
+def test_is_conditional_dependent_static_shape():
+    """Test that we don't consider dependencies through "constant" shape Ops"""
+    x1 = pt.matrix("x1", shape=(None, 5))
+    y1 = pt.random.normal(size=pt.shape(x1))
+    assert is_conditional_dependent(y1, x1, [x1, y1])
+
+    x2 = pt.matrix("x2", shape=(9, 5))
+    y2 = pt.random.normal(size=pt.shape(x2))
+    assert not is_conditional_dependent(y2, x2, [x2, y2])