Fix issue with sampling of PartialObservedRVs

pymc/distributions/distribution.py

@@ -1337,7 +1337,8 @@ def create_partial_observed_rv(
             ndim_supp=rv.owner.op.ndim_supp,
         )(rv, mask)
 
-    joined_rv = pt.empty(rv.shape, dtype=rv.type.dtype)
+    [rv_shape] = constant_fold([rv.shape], raise_not_constant=False)
+    joined_rv = pt.empty(rv_shape, dtype=rv.type.dtype)
     joined_rv = pt.set_subtensor(joined_rv[mask], unobserved_rv)
     joined_rv = pt.set_subtensor(joined_rv[antimask], observed_rv)
 

pymc/pytensorf.py

@@ -56,7 +56,6 @@
     RandomGeneratorSharedVariable,
     RandomStateSharedVariable,
 )
-from pytensor.tensor.rewriting.basic import topo_constant_folding
 from pytensor.tensor.rewriting.shape import ShapeFeature
 from pytensor.tensor.sharedvar import SharedVariable, TensorSharedVariable
 from pytensor.tensor.subtensor import AdvancedIncSubtensor, AdvancedIncSubtensor1
@@ -1015,7 +1014,8 @@ def constant_fold(
     """
     fg = FunctionGraph(outputs=xs, features=[ShapeFeature()], clone=True)
 
-    folded_xs = rewrite_graph(fg, custom_rewrite=topo_constant_folding).outputs
+    # By default, rewrite_graph includes canonicalize which includes constant-folding as the final rewrite
+    folded_xs = rewrite_graph(fg).outputs
 
     if raise_not_constant and not all(isinstance(folded_x, Constant) for folded_x in folded_xs):
         raise NotConstantValueError

tests/distributions/test_distribution.py

@@ -1061,3 +1061,17 @@ def test_wrong_mask(self):
         invalid_mask = np.zeros((1, 5), dtype=bool)
         with pytest.raises(ValueError, match="mask can't have more dims than rv"):
             create_partial_observed_rv(rv, invalid_mask)
+
+    @pytest.mark.filterwarnings("error")
+    def test_default_updates(self):
+        mask = np.array([True, True, False])
+        rv = pm.Normal.dist(shape=(3,))
+        (obs_rv, _), (unobs_rv, _), joined_rv = create_partial_observed_rv(rv, mask)
+
+        draws_obs_rv, draws_unobs_rv, draws_joined_rv = pm.draw(
+            [obs_rv, unobs_rv, joined_rv], draws=2
+        )
+
+        assert np.all(draws_obs_rv[0] != draws_obs_rv[1])
+        assert np.all(draws_unobs_rv[0] != draws_unobs_rv[1])
+        assert np.all(draws_joined_rv[0] != draws_joined_rv[1])

tests/sampling/test_jax.py

@@ -30,6 +30,7 @@
 
 import pymc as pm
 
+from pymc import ImputationWarning
 from pymc.distributions.multivariate import PosDefMatrix
 from pymc.sampling.jax import (
     _get_batched_jittered_initial_points,
@@ -459,3 +460,14 @@ def test_idata_contains_stats(sampler_name: str):
     for stat_var, stat_var_dims in stat_vars.items():
         assert stat_var in stats.variables
         assert stats.get(stat_var).values.shape == stat_var_dims
+
+
+def test_sample_partially_observed():
+    with pm.Model() as m:
+        with pytest.warns(ImputationWarning):
+            x = pm.Normal("x", observed=np.array([0, 1, np.nan]))
+        idata = pm.sample(nuts_sampler="numpyro", chains=1, draws=10, tune=10)
+
+    assert idata.observed_data["x_observed"].shape == (2,)
+    assert idata.posterior["x_unobserved"].shape == (1, 10, 1)
+    assert idata.posterior["x"].shape == (1, 10, 3)