diff --git a/pymc/backends/arviz.py b/pymc/backends/arviz.py index 05043ca6ff..1b82ddf3a3 100644 --- a/pymc/backends/arviz.py +++ b/pymc/backends/arviz.py @@ -153,15 +153,21 @@ def find_observations(model: Model) -> dict[str, Var]: def find_constants(model: Model) -> dict[str, Var]: """If there are constants available, return them as a dictionary.""" - model_vars = model.basic_RVs + model.deterministics + model.potentials - value_vars = set(model.rvs_to_values.values()) + # Collect data_vars that feed into observed value variables. + # This handles both the direct case (value_var IS the data var, e.g. continuous + # distributions) and the indirect case (value_var wraps the data var through a + # Cast op, e.g. discrete distributions like Categorical or Poisson). + observed_data_vars = set() + for rv in model.observed_RVs: + value_var = model.rvs_to_values[rv] + for anc in ancestors([value_var]): + if anc in model.data_vars: + observed_data_vars.add(anc) constant_data = {} for var in model.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 + if var in observed_data_vars: + continue if isinstance(var, SharedVariable): var_value = var.get_value() diff --git a/tests/backends/test_arviz.py b/tests/backends/test_arviz.py index 6211c6bf9c..d6d7a3c79b 100644 --- a/tests/backends/test_arviz.py +++ b/tests/backends/test_arviz.py @@ -480,7 +480,7 @@ def test_constant_data(self, use_context): @pytest.mark.parametrize("constant_in_generative_graph", [True, False]) def test_observed_data_also_constant(self, constant_in_generative_graph): - """Test that wen the same variable is used as constant data and observed data, it shows up in both groups.""" + """Test that observed data never leaks into constant_data, even when also used in the generative graph.""" with pm.Model(coords={"trial": [0, 1, 2]}) as model: x = pm.Data("x", [1.0, 2.0, 3.0], dims=["trial"]) sigma = pm.HalfNormal("sigma", 1) @@ -494,11 +494,28 @@ def test_observed_data_also_constant(self, constant_in_generative_graph): test_dict = { "prior": ["sigma"], "observed_data": ["y"], + "~constant_data": [], + } + fails = check_multiple_attrs(test_dict, inference_data) + assert not fails + + def test_discrete_observed_not_in_constant_data(self): + """Regression test for #7851: discrete observed data should not leak into constant_data.""" + with pm.Model() as model: + y_obs = pm.Data("y_obs", [0, 1, 2]) + beta = pm.Normal("beta", shape=3) + p = pm.math.softmax(beta) + pm.Categorical("y", p=p, observed=y_obs) + + trace = pm.sample_prior_predictive(100, return_inferencedata=False) + + inference_data = to_inference_data(prior=trace, model=model, log_likelihood=False) + + test_dict = { + "prior": ["beta"], + "observed_data": ["y"], + "~constant_data": [], } - if constant_in_generative_graph: - test_dict["constant_data"] = ["x"] - else: - test_dict["~constant_data"] = [] fails = check_multiple_attrs(test_dict, inference_data) assert not fails