Fix find_constants excluding discrete observed data from constant_data

[genrichez]

Closes #7851

Problem

When using discrete distributions (Categorical, Poisson, Binomial) with pm.Data as observed, the data variable incorrectly ends up in the constant_data group of InferenceData. It should only appear in observed_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)

idata = pm.sample_prior_predictive()
# y_obs incorrectly shows up in idata.constant_data

Cause

For discrete distributions, the observed value variable gets wrapped in a Cast op (e.g. Cast{int64}(y_obs)). The old code in find_constants checked var in value_vars, but value_vars contains the Cast output, not the original pm.Data variable. So the identity check fails and the data var slips through into constant_data.

For continuous distributions this worked fine because the value variable IS the data var directly (no Cast in between).

Fix

Instead of checking if a data var is directly in value_vars, walk the ancestors of each observed RV's value variable and collect any data_vars found there. This catches the data var regardless of whether a Cast (or any other transform) sits between it and the value var.

Tests

• Added test_discrete_observed_not_in_constant_data as a regression test for the reported bug
• Updated test_observed_data_also_constant to reflect that observed data should not appear in constant_data even when also used in the generative graph (per discussion in BUG: find_constants helper function is inconsistent #7851)
• All 50 tests in tests/backends/test_arviz.py pass

[read-the-docs-community]

Documentation build overview

📚 pymc | 🛠️ Build #33632418 | 📁 Comparing 4b6e28d against latest (df1a369)

  🔍 Preview build  

2 files changed

± glossary.html
± _modules/pymc/backends/arviz.html

[williambdean]

Why the return_inferencedata=False. Does this occur with True?

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 88.16%. Comparing base (df1a369) to head (4b6e28d).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #8368      +/-   ##
==========================================
- Coverage   91.79%   88.16%   -3.64%     
==========================================
  Files         128      128              
  Lines       21007    21010       +3     
==========================================
- Hits        19283    18523     -760     
- Misses       1724     2487     +763     

Files with missing lines	Coverage Δ
pymc/backends/arviz.py	93.50% <100.00%> (-2.89%)	⬇️

... and 18 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.