PyMC v4 not returning 'observed_data' in InferenceData object

Questions v5
4

Sorry it took so long to get this to you - it looks like they closed the ticket.

Please see the following simulated data:
process_data.csv (103.8 KB)

# Test treatments
tot_count = df['value'].values

process_idx, process = df['process'].factorize(sort=True)

COORDS = {

    'obs_id': np.arange(len(process_idx)),
    'process': process

}

with pm.Model(coords = COORDS) as example_model:

    ###
    # Test treatments
    # Data
    count = pm.Data('total_count', tot_count, dims='obs_id')

    # hyper Exponential
    mu_lam = pm.Exponential('mu_lam', lam=0.3)

    # hyper Exponential
    a_lam = pm.Exponential('a_lam', lam=2)

    # Prior - mu
    mu = pm.Exponential('mu', lam= mu_lam, dims=('process'))
    
    # Prior alpha
    alpha = pm.Exponential('alpha', lam= a_lam, dims=('process'))


    # Treatment Likelihoods
    obs = pm.NegativeBinomial('obs', mu= mu[process_idx], alpha= alpha[process_idx], observed= count, dims= ('obs_id'))
    
    # TRACE

    idata_trace = pm.sampling_jax.sample_numpyro_nuts(tune=3_000, draws=5_000, target_accept=0.95)

    idata_trace.extend(
        pm.sample_posterior_predictive(idata_trace, \
        var_names=['mu', 'alpha', 'obs']
                      )
    )

    idata_trace.extend(
        pm.sample_prior_predictive(1_000)
    )

The trace returned does not include ‘observed_data’:

image

There is this warning as well:

/home/user/anaconda3/envs/pymc-dev-py39-gpu/lib/python3.9/site-packages/pymc/backends/arviz.py:58: UserWarning: Could not extract data from symbolic observation obs
  warnings.warn(f"Could not extract data from symbolic observation {obs}")