Hierarchical Partially Pooled Media

[Msettnes]

trafficstars

I'm looking to make a partially pooled media variable. Such that Facebook and TV ads can not only get separate priors, but be pooled such that the individual coefficient for media will be media_coef_tv_i = N(mu_tv, sigma_tv) media_coef_fb_i = N(mu_fb, sigma_fb)

This stabilises the media contribution by pooling parameters allowing for greater granularity (when reduced granularity leads to less data), while allowing for different type of media to show different behaviour.

I have done such models before directly in STAN, but would like to modify LMMM in a way so I can still take advantage of the budget optimizer?

Can someone more experienced with numpyro guide me in the right direction?