How to partially pool within subcategory

Questions version agnostic
1

I currently have a model that pools a parameter for a region regardless of what country that region is in:

region_sigma = pm.Exponential("sigma", 1)
region_intercept = pm.Normal("region_intercept", mu=0, sigma=region_sigma, dims="region_id_dims")

But I want to change this so that the sigma for regions is pooled only for those regions within the same country. How would I go about doing this in PYMC4, please? I’ve tried the following:

region_sigma = pm.Exponential("region_sigma", 1, dims="country_dims")
region_intercept = pm.Normal("region_intercept", mu=0, sigma=region_sigma[country_idx], dims="region_dims")

Where country_idx is a numerical index variable for each row of data, corresponding to the country. But I get the error:

ValueError: conflicting sizes for dimensions 'region_id_dims'

Thanks!

2

Can you give us a sense of the shape of your country/region indices? For example, I’m not clear what a “region” is or how regions map to countries.