BUG: `find_constants` helper function is inconsistent

[OriolAbril]

Describe the issue:

The function find_constants in backends/arviz.py aims to return the elements that should go into the constant data group and has some logic to remove instances of pm.Data that are later on used as observed= so they don't end up duplicated in both observed_data and constant_data.

This exclusion doesn't always seem to work though. With pm.Categorical as observational distribution, the observed data ends up duplicated in both observed and constant data groups.

Reproduceable code example:

import numpy as np
import pymc as pm

rng = np.random.default_rng(1234)

obs_value = rng.choice(list("ABCDE"), 100, replace=True)

groups, groups_idx = np.unique(obs_value, return_inverse=True)

coords = {
    "obs_idx": np.arange(100),
    "groups": groups,
}

with pm.Model(coords=coords) as model:
    y_pt = pm.Data("y_observed", groups_idx, dims="obs_idx")
    beta = pm.Normal("beta", dims="groups")
    p = pm.Deterministic("p", pm.math.softmax(beta), dims="groups")
    pm.Categorical("y", p=p, observed=y_pt, dims="obs_idx")

with model:
    idata = pm.sample()

np.allclose(idata.constant_data["y_observed"], idata.observed_data["y"])
# True

Error message:

PyMC version information:

PyMC: 5.23.0 installed from conda.

Not sure if this is something that has always happened or if it is something new as I generally don't register the observed values as pm.Data.

Context for the issue:

This means, for example, that in more complex scenarios where we have both actual constant data and the observed data it is not possible to update the pm.Data containers except the y_observed to something with a different length and call pm.sample_posterior_predictive. As y_observed is still being pulled in, the conversion to inferencedata after sampling the posterior predictive fails because of incompatible shapes.

This seems to come from the set(model.rvs_to_values.values()) returning something related to casting y_observed instead of y_observed itself.

# for the model above:
set(model.rvs_to_values.values())
# {Cast{int64}.0, beta}

# whereas for a model with Normal as observational distribution
with pm.Model(coords={"obs_idx": np.arange(100)}) as model:
    y_pt = pm.Data("y_observed", rng.normal(size=100), dims="obs_idx")
    sigma = pm.Normal("sigma")
    pm.Normal("y", mu=0, sigma=sigma, observed=y_pt)

set(model.rvs_to_values.values())
# {sigma, y_observed}

# another discrete model has the same issue
with pm.Model(coords={"obs_idx": np.arange(100)}) as model:
    y_pt = pm.Data("y_observed", rng.poisson(size=100), dims="obs_idx")
    lam = pm.Normal("lambda")
    pm.Poisson("y", lam, observed=y_pt)

set(model.rvs_to_values.values())
# {Cast{int64}.0, lambda}

[ricardoV94]

This was the logic that was supposed to filter away observed variables from data_vars:

        if var in value_vars:
            # An observed value variable could also be part of the generative graph
            if var not in ancestors(model_vars):
                continue

I guess we could check all the ancestors of value_vars and be sure var is not there:

value_vars_ancestors = set(ancestors(value_vars))
for var in model.data_vars:
    if var in value_vars_ancestors:
        continue

What does arviz want to do if the value_variable is part of the generative graph as well? Weird but valid case.

with pm.Model() as m:
  x = pm.Data("x")
  mu = pm.Normal("beta") * x
  y = pm.Normal("y", mu, observed=x)

Right now that's what we do with var not in ancestors(model_vars) second check before continue. We can do the same, in which case we can probably cache the difference as value_vars_ancestors = set(ancestors(value_vars)) - set(ancestors(model_vars)))

[OriolAbril]

I think in the last example x should still go only into observed_data. For all the data the model is conditioned on, the distinction between observed/constant is basically "is modeled as a random variable" -> observed data, otherwise constant data. I don't think it matters if in addition to being modelled as a random variable we are also using it somewhere else in the model.

[ricardoV94]

Sure. Do you want to open a PR?

[OriolAbril]

I will put it on the list, but being realistic I won't get to this very soon. If someone else wants to send the PR don't hesitate to do it on my account.