Repository: pymc-devs/pymc (Python)

Issue:
BUG: NotConstantValueError when using coordinates observed data with missing values with pymc==5.14.0
### Describe the issue:

In pymc 5.14.0, I get a `NotConstantValueError` when sampling with observed data that contains missing values. This only occurs when passing coords to the model.

The same code runs without an error with pymc==5.12.0.

### Reproduceable code example:

```python
import numpy as np
import pymc as pm
size = 30
f1 = np.random.normal(scale=1.5, size=size)
f2 = np.random.normal(scale=0.5, size=size)
features = np.stack((f1, f2), axis=1)
features[0, 0] = np.nan
n_features = features.shape[1]

coords = {"location": range(size),
          "feature": ("f1", "f2")}

with pm.Model(coords=coords) as model:
    sd_dist = pm.Exponential.dist(1.0, shape=n_features)
    chol, _, _ = pm.LKJCholeskyCov('chol_cov', n=n_features, eta=2,
                                   sd_dist=sd_dist, compute_corr=True)
    vals = pm.MvNormal('vals', mu=np.zeros(n_features), chol=chol,
                       observed=features, dims=("location", "feature"))
    trace = pm.sample()
```


### Error message:

```shell
/Users/jennifer/env/lib/python3.11/site-packages/pymc/model/core.py:1349: ImputationWarning: Data in vals contains missing values and will be automatically imputed from the sampling distribution.
  warnings.warn(impute_message, ImputationWarning)
---------------------------------------------------------------------------
NotConstantValueError                     Traceback (most recent call last)
Cell In[2], line 19
     15 chol, _, _ = pm.LKJCholeskyCov('chol_cov', n=n_features, eta=2,
     16                                sd_dist=sd_dist, compute_corr=True)
     17 vals = pm.MvNormal('vals', mu=np.zeros(n_features), chol=chol,
     18                    observed=features, dims=("location", "feature"))
---> 19 trace = pm.sample()

File ~/env/lib/python3.11/site-packages/pymc/sampling/mcmc.py:684, in sample(draws, tune, chains, cores, random_seed, progressbar, progressbar_theme, step, var_names, nuts_sampler, initvals, init, jitter_max_retries, n_init, trace, discard_tuned_samples, compute_convergence_checks, keep_warning_stat, return_inferencedata, idata_kwargs, nuts_sampler_kwargs, callback, mp_ctx, model, **kwargs)
    681         auto_nuts_init = False
    683 initial_points = None
--> 684 step = assign_step_methods(model, step, methods=pm.STEP_METHODS, step_kwargs=kwargs)
    686 if nuts_sampler != "pymc":
    687     if not isinstance(step, NUTS):

File ~/env/lib/python3.11/site-packages/pymc/sampling/mcmc.py:212, in assign_step_methods(model, step, methods, step_kwargs)
    210 methods_list: list[type[BlockedStep]] = list(methods or pm.STEP_METHODS)
    211 selected_steps: dict[type[BlockedStep], list] = {}
--> 212 model_logp = model.logp()
    214 for var in model.value_vars:
    215     if var not in assigned_vars:
    216         # determine if a gradient can be computed

File ~/env/lib/python3.11/site-packages/pymc/model/core.py:725, in Model.logp(self, vars, jacobian, sum)
    723 rv_logps: list[TensorVariable] = []
    724 if rvs:
--> 725     rv_logps = transformed_conditional_logp(
    726         rvs=rvs,
    727         rvs_to_values=self.rvs_to_values,
    728         rvs_to_transforms=self.rvs_to_transforms,
    729         jacobian=jacobian,
    730     )
    731     assert isinstance(rv_logps, list)
    733 # Replace random variables by their value variables in potential terms

File ~/env/lib/python3.11/site-packages/pymc/logprob/basic.py:611, in transformed_conditional_logp(rvs, rvs_to_values, rvs_to_transforms, jacobian, **kwargs)
    608     transform_rewrite = TransformValuesRewrite(values_to_transforms)  # type: ignore
    610 kwargs.setdefault("warn_rvs", False)
--> 611 temp_logp_terms = conditional_logp(
    612     rvs_to_values,
    613     extra_rewrites=transform_rewrite,
    614     use_jacobian=jacobian,
    615     **kwargs,
    616 )
    618 # The function returns the logp for every single value term we provided to it.
    619 # This includes the extra values we plugged in above, so we filter those we
    620 # actually wanted in the same order they were given in.
    621 logp_terms = {}

File ~/env/lib/python3.11/site-packages/pymc/logprob/basic.py:541, in conditional_logp(rv_values, warn_rvs, ir_rewriter, extra_rewrites, **kwargs)
    538 q_values = remapped_vars[: len(q_values)]
    539 q_rv_inputs = remapped_vars[len(q_values) :]
--> 541 q_logprob_vars = _logprob(
    542     node.op,
    543     q_values,
    544     *q_rv_inputs,
    545     **kwargs,
    546 )
    548 if not isinstance(q_logprob_vars, list | tuple):
    549     q_logprob_vars = [q_logprob_vars]

File /opt/local/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/functools.py:909, in singledispatch.<locals>.wrapper(*args, **kw)
    905 if not args:
    906     raise TypeError(f'{funcname} requires at least '
    907                     '1 positional argument')
--> 909 return dispatch(args[0].__class__)(*args, **kw)

File ~/env/lib/python3.11/site-packages/pymc/distributions/distribution.py:1633, in partial_observed_rv_logprob(op, values, dist, mask, **kwargs)
   1631 [obs_value, unobs_value] = values
   1632 antimask = ~mask
-> 1633 joined_value = pt.empty(constant_fold([dist.shape])[0])
   1634 joined_value = pt.set_subtensor(joined_value[mask], unobs_value)
   1635 joined_value = pt.set_subtensor(joined_value[antimask], obs_value)

File ~/env/lib/python3.11/site-packages/pymc/pytensorf.py:1037, in constant_fold(xs, raise_not_constant)
   1034 folded_xs = rewrite_graph(fg).outputs
   1036 if raise_not_constant and not all(isinstance(folded_x, Constant) for folded_x in folded_xs):
-> 1037     raise NotConstantValueError
   1039 return tuple(
   1040     folded_x.data if isinstance(folded_x, Constant) else folded_x for folded_x in folded_xs
   1041 )

NotConstantValueError:
```


### PyMC version information:

pymc==5.14.0
pytensor==2.20.0

MacOS Sonoma 14.4.1
Installation: pip

### Context for the issue:

Using PyMC for data analysis for a customer. Covariates for certain locations have missing values, but still need to provide predictions for these locations to the customer. Passing in coordinates is important to keep track of different locations in the trace for visualization.

Currently working around by using pymc==5.12.0.


Discussion:
![Welcome Banner](https://raw.githubusercontent.com/pymc-devs/brand/main/welcome-bot/BannerWelcome.jpg)]
:tada: Welcome to _PyMC_! :tada: We're really excited to have your input into the project! :sparkling_heart:
<br>If you haven't done so already, please make sure you check out our [Contributing Guidelines](https://www.pymc.io/projects/docs/en/latest/contributing/index.html) and [Code of Conduct](https://github.com/pymc-devs/pymc/blob/main/CODE_OF_CONDUCT.md).

Thanks, we can try to pass the flag `raise_not_constant=False` when we call `constant_fold` there
Should be fixed by #7314 



Patch:
diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py
index 214ec505d7..bcd4613918 100644
--- a/pymc/distributions/distribution.py
+++ b/pymc/distributions/distribution.py
@@ -1630,7 +1630,9 @@ def partial_observed_rv_logprob(op, values, dist, mask, **kwargs):
     # For the logp, simply join the values
     [obs_value, unobs_value] = values
     antimask = ~mask
-    joined_value = pt.empty(constant_fold([dist.shape])[0])
+    # We don't need it to be completely folded, just to avoid any RVs in the graph of the shape
+    [folded_shape] = constant_fold([dist.shape], raise_not_constant=False)
+    joined_value = pt.empty(folded_shape)
     joined_value = pt.set_subtensor(joined_value[mask], unobs_value)
     joined_value = pt.set_subtensor(joined_value[antimask], obs_value)
     joined_logp = logp(dist, joined_value)


Test patch:
diff --git a/tests/distributions/test_distribution.py b/tests/distributions/test_distribution.py
index d5f3359dd1..caac1777bc 100644
--- a/tests/distributions/test_distribution.py
+++ b/tests/distributions/test_distribution.py
@@ -979,8 +979,9 @@ def test_univariate(self, symbolic_rv):
         np.testing.assert_allclose(obs_logp, st.norm([1, 2]).logpdf([0.25, 0.5]))
         np.testing.assert_allclose(unobs_logp, st.norm([3]).logpdf([0.25]))
 
+    @pytest.mark.parametrize("mutable_shape", (False, True))
     @pytest.mark.parametrize("obs_component_selected", (True, False))
-    def test_multivariate_constant_mask_separable(self, obs_component_selected):
+    def test_multivariate_constant_mask_separable(self, obs_component_selected, mutable_shape):
         if obs_component_selected:
             mask = np.zeros((1, 4), dtype=bool)
         else:
@@ -988,7 +989,11 @@ def test_multivariate_constant_mask_separable(self, obs_component_selected):
         obs_data = np.array([[0.1, 0.4, 0.1, 0.4]])
         unobs_data = np.array([[0.4, 0.1, 0.4, 0.1]])
 
-        rv = pm.Dirichlet.dist([1, 2, 3, 4], shape=(1, 4))
+        if mutable_shape:
+            shape = (1, pytensor.shared(np.array(4, dtype=int)))
+        else:
+            shape = (1, 4)
+        rv = pm.Dirichlet.dist(pt.arange(shape[-1]) + 1, shape=shape)
         (obs_rv, obs_mask), (unobs_rv, unobs_mask), joined_rv = create_partial_observed_rv(rv, mask)
 
         # Test types
@@ -1023,6 +1028,10 @@ def test_multivariate_constant_mask_separable(self, obs_component_selected):
         np.testing.assert_allclose(obs_logp, expected_obs_logp)
         np.testing.assert_allclose(unobs_logp, expected_unobs_logp)
 
+        if mutable_shape:
+            shape[-1].set_value(7)
+            assert tuple(joined_rv.shape.eval()) == (1, 7)
+
     def test_multivariate_constant_mask_unseparable(self):
         mask = pt.constant(np.array([[True, True, False, False]]))
         obs_data = np.array([[0.1, 0.4, 0.1, 0.4]])
@@ -1097,14 +1106,19 @@ def test_multivariate_shared_mask_separable(self):
         np.testing.assert_almost_equal(obs_logp, new_expected_logp)
         np.testing.assert_array_equal(unobs_logp, [])
 
-    def test_multivariate_shared_mask_unseparable(self):
+    @pytest.mark.parametrize("mutable_shape", (False, True))
+    def test_multivariate_shared_mask_unseparable(self, mutable_shape):
         # Even if the mask is initially not mixing support dims,
         # it could later be changed in a way that does!
         mask = shared(np.array([[True, True, True, True]]))
         obs_data = np.array([[0.1, 0.4, 0.1, 0.4]])
         unobs_data = np.array([[0.4, 0.1, 0.4, 0.1]])
 
-        rv = pm.Dirichlet.dist([1, 2, 3, 4], shape=(1, 4))
+        if mutable_shape:
+            shape = mask.shape
+        else:
+            shape = (1, 4)
+        rv = pm.Dirichlet.dist([1, 2, 3, 4], shape=shape)
         (obs_rv, obs_mask), (unobs_rv, unobs_mask), joined_rv = create_partial_observed_rv(rv, mask)
 
         # Test types
@@ -1134,16 +1148,22 @@ def test_multivariate_shared_mask_unseparable(self):
 
         # Test that we can update a shared mask
         mask.set_value(np.array([[False, False, True, True]]))
+        equivalent_value = np.array([0.1, 0.4, 0.4, 0.1])
 
         assert tuple(obs_rv.shape.eval()) == (2,)
         assert tuple(unobs_rv.shape.eval()) == (2,)
 
-        new_expected_logp = pm.logp(rv, [0.1, 0.4, 0.4, 0.1]).eval()
+        new_expected_logp = pm.logp(rv, equivalent_value).eval()
         assert not np.isclose(expected_logp, new_expected_logp)  # Otherwise test is weak
         obs_logp, unobs_logp = logp_fn()
         np.testing.assert_almost_equal(obs_logp, new_expected_logp)
         np.testing.assert_array_equal(unobs_logp, [])
 
+        if mutable_shape:
+            mask.set_value(np.array([[False, False, True, False], [False, False, False, True]]))
+            assert tuple(obs_rv.shape.eval()) == (6,)
+            assert tuple(unobs_rv.shape.eval()) == (2,)
+
     def test_support_point(self):
         x = pm.GaussianRandomWalk.dist(init_dist=pm.Normal.dist(-5), mu=1, steps=9)
         ref_support_point = support_point(x).eval()
diff --git a/tests/model/test_core.py b/tests/model/test_core.py
index cdb93c72a8..c00250b739 100644
--- a/tests/model/test_core.py
+++ b/tests/model/test_core.py
@@ -28,6 +28,7 @@
 import pytensor.sparse as sparse
 import pytensor.tensor as pt
 import pytest
+import scipy
 import scipy.sparse as sps
 import scipy.stats as st
 
@@ -38,7 +39,7 @@
 
 import pymc as pm
 
-from pymc import Deterministic, Model, Potential
+from pymc import Deterministic, Model, MvNormal, Potential
 from pymc.blocking import DictToArrayBijection, RaveledVars
 from pymc.distributions import Normal, transforms
 from pymc.distributions.distribution import PartialObservedRV
@@ -1504,11 +1505,39 @@ def test_truncated_normal(self):
         """
         with Model() as m:
             mu = pm.TruncatedNormal("mu", mu=1, sigma=2, lower=0)
-            x = pm.TruncatedNormal(
-                "x", mu=mu, sigma=0.5, lower=0, observed=np.array([0.1, 0.2, 0.5, np.nan, np.nan])
-            )
+            with pytest.warns(ImputationWarning):
+                x = pm.TruncatedNormal(
+                    "x",
+                    mu=mu,
+                    sigma=0.5,
+                    lower=0,
+                    observed=np.array([0.1, 0.2, 0.5, np.nan, np.nan]),
+                )
         m.check_start_vals(m.initial_point())
 
+    def test_coordinates(self):
+        # Regression test for https://github.com/pymc-devs/pymc/issues/7304
+
+        coords = {"trial": range(30), "feature": range(2)}
+        observed = np.zeros((30, 2))
+        observed[0, 0] = np.nan
+
+        with Model(coords=coords) as model:
+            with pytest.warns(ImputationWarning):
+                MvNormal(
+                    "y",
+                    mu=np.zeros(2),
+                    cov=np.eye(2),
+                    observed=observed,
+                    dims=("trial", "feature"),
+                )
+
+        logp_fn = model.compile_logp()
+        np.testing.assert_allclose(
+            logp_fn({"y_unobserved": [0]}),
+            scipy.stats.multivariate_normal.logpdf([0, 0], cov=np.eye(2)) * 30,
+        )
+
 
 class TestShared:
     def test_deterministic(self):