Change rng of partially observed RVs

Author: ricardoV94

pymc/model.py

@@ -1276,6 +1276,16 @@ def make_obs_var(
                 clone=False,
             )
             (observed_rv_var,) = local_subtensor_rv_lift.transform(fgraph, fgraph.outputs[0].owner)
+            # Make a clone of the RV, but change the rng so that observed and missing
+            # are not treated as equivalent nodes by aesara. This would happen if the
+            # size of the masked and unmasked array happened to coincide
+            _, size, _, *inps = observed_rv_var.owner.inputs
+            rng = self.model.next_rng()
+            observed_rv_var = observed_rv_var.owner.op(*inps, size=size, rng=rng)
+            # Add default_update to new rng
+            new_rng = observed_rv_var.owner.outputs[0]
+            observed_rv_var.update = (rng, new_rng)
+            rng.default_update = new_rng
             observed_rv_var.name = f"{name}_observed"
 
             observed_rv_var.tag.observations = nonmissing_data

pymc/tests/test_missing.py

@@ -18,8 +18,10 @@
 import pytest
 import scipy.stats
 
+from aesara.graph import graph_inputs
 from numpy import array, ma
 
+from pymc import logpt
 from pymc.distributions import Dirichlet, Gamma, Normal, Uniform
 from pymc.exceptions import ImputationWarning
 from pymc.model import Model
@@ -207,3 +209,26 @@ def test_missing_vector_parameter():
         m.logp({"x_missing": np.array([-10, 10, -10, 10])}),
         scipy.stats.norm(scale=0.1).logpdf(0) * 6,
     )
+
+
+def test_missing_symmetric():
+    """Check that logpt works when partially observed variable have equal observed and
+    unobserved dimensions.
+
+    This would fail in a previous implementation because the two variables would be
+    equivalent and one of them would be discarded during MergeOptimization while
+    buling the logpt graph
+    """
+    with Model() as m:
+        x = Gamma("x", alpha=3, beta=10, observed=np.array([1, np.nan]))
+
+    x_obs_rv = m["x_observed"]
+    x_obs_vv = m.rvs_to_values[x_obs_rv]
+
+    x_unobs_rv = m["x_missing"]
+    x_unobs_vv = m.rvs_to_values[x_unobs_rv]
+
+    logp = logpt([x_obs_rv, x_unobs_rv], {x_obs_rv: x_obs_vv, x_unobs_rv: x_unobs_vv})
+    logp_inputs = list(graph_inputs([logp]))
+    assert x_obs_vv in logp_inputs
+    assert x_unobs_vv in logp_inputs