I have a simple multi-level model without any parameter pooling in the model defined in the following way without any pooling of parameters:

```
coords = {"Country":country_list, "obs_id":np.arange(yearly.shape[0])}
with pm.Model(coords=coords) as unpooled_model:
country_idx = pm.Data("country_idx", yearly.country_idx, dims="obs_id")
a = pm.Exponential("a", 1/10, dims="Country")
counts = pm.Poisson("counts", a[yearly.country_idx], observed = yearly["Gold"], dims="obs_id")
```

I can obtain posterior predictive samples from this model straightforwardly and visualise it with arviz. However, each country index (country_idx) value appears a different number of times within this dataset, and this results in getting imbalanced posterior predictive samples, coming from each country. What I would really like to do is only get posterior predictive samples for certain country_idx values. Whenever I have tried to do this using pm.set_data({“country_idx”:[1.4]}) for instance, I get errors regarding conflicts with the shape of the obs_id dimension in the original model. Is there a smart/easy way to do this?

I have looked through the tutorials on multi-level models and this isn’t done there either. For instance, in the linked tutorial, imagine I wanted posterior predictive samples, but only for the aitkin basement, and not the others. What is the best way to do this in pymc3?

Thanks in advance