Multilevel hierarchical modelling with labels - how to set data for prediction

Questions v5
,
1

Hi All,
I am trying to build a hierarchical bayesian model for sales across different countries. I need some flexibility in the naming of the columns since I need the model to be configurable which is the reason behind the dimension for the media variables.
This model works and produces decent results but I don’t know how to pm.set_data for predictions in this case due to the loop in the model definition.
For context the data has the following structure:

Date Country Sales TV Radio Print
0 Germany 4
1 Germany 5
2 Germany 4
0 Austria 10
1 Austria 12
2 Austria 13

And The model is the following:

import pymc as pm
import pytensor.tensor as pt

def geometric_adstock(x, alpha: float = 0.0, l_max: int = 12):
    """Geometric adstock transformation."""
    cycles = [
        pt.concatenate(
            [pt.zeros([i, x.shape[1]]), x[: x.shape[0] - i]]
        )
        for i in range(l_max)
    ]
    x_cycle = pt.stack(cycles, axis=1)
    w_1 = pt.as_tensor_variable([pt.power(alpha, i) for i in range(l_max)])
    return pt.sum(pt.mul(w_1, x_cycle),axis=1)

def logistic_saturation(x, lam: float = 0.5):
    """Logistic saturation transformation."""
    return (1 - pt.exp(-lam * x)) / (1 + pt.exp(-lam * x))

coords = {
    "media_names": X[['TV', 'Radio', 'Banners']].columns.values,
    "country_names": ['Germany', 'Austria'],
    "date": X['date'].unique(),
    "obs_id": X.index.values
}

with pm.Model(coords= coords) as bhmmm:
    # Hyperpriors
    coef_lam = pm.Exponential("coef_lam", lam=10, dims = "media_names")
    sat_lam = pm.Exponential("sat_lam", lam=10, dims = "media_names")
    car_alpha = pm.Exponential("car_alpha", lam=0.01, dims= "media_names")
    car_beta = pm.Exponential("car_beta", lam=0.01, dims = "media_names")
    base_lam = pm.Exponential("base_lam", lam=10)

    # For each country
    for country in X["Country"].unique():
        X_ = X[X["Country"] == country]

        media = pm.MutableData(
            name=f"media_{country}", 
            value =  X_[["TV", "Radio", "Banners"]],
            dims = ('date', 'media_names'))
        
        coef = pm.Exponential(f"coef_{country}", lam=coef_lam, dims = ("media_names"))
        sat = pm.Exponential(f"sat_{country}", lam=sat_lam, dims = ("media_names"))
        car = pm.Beta(f"car_{country}",  alpha=car_alpha, beta=car_beta, dims= ("media_names"))

        media_adstock = pm.Deterministic(
            name=f"media_adstock_{country}", 
            var=geometric_adstock(x=media, alpha=car, l_max=12), 
            dims=("date", "media_names"))
        media_adstock_saturated = pm.Deterministic(
            name=f"media_adstock_saturated_{country}",
            var=logistic_saturation(x=media_adstock, lam = sat),
            dims=("date", "media_names"))
        regression_media = pm.Deterministic(
            name=f"regression_media_{country}", 
            var= media_adstock_saturated * coef, 
            dims = ('date', 'media_names'))

        base = pm.Exponential(f"base_{country}", lam=base_lam)
        noise = pm.Exponential(f"noise_{country}", lam=0.001)

        sales = pm.Normal(
            f"sales_{country}",
            mu=regression_media.sum(axis=-1) + base,
            sigma=noise,
            observed=y[X_.index].values,
        )

    trace = pm.sample(tune=3000)

I have also tried to rewrite the code following the very useful [radon example] (A Primer on Bayesian Methods for Multilevel Modeling — PyMC example gallery) but I am falling short on how to integrate the media dimensionality. Here is what I wrote so far.

with pm.Model(coords= coords) as bhmmm:
    # Hyperpriors
    coef_lam = pm.Exponential("coef_lam", lam=10, dims = "media_names")
    sat_lam = pm.Exponential("sat_lam", lam=10, dims = "media_names")
    car_alpha = pm.Exponential("car_alpha", lam=0.01, dims= "media_names")
    car_beta = pm.Exponential("car_beta", lam=0.01, dims = "media_names")
    base_lam = pm.Exponential("base_lam", lam=10)

    country_idx = pm.MutableData(name="country", 
                             value = country, 
                             dims="obs_id")
    media = pm.MutableData(
            name="media", 
            value =  X[["TV", "Radio", "Banners"]],
            dims = ('obs_id', 'media_names'))
    
    coef = pm.Exponential("coef", lam=coef_lam, dims = ("media_names","country_names"))
    sat = pm.Exponential("sat", lam=sat_lam, dims = ("media_names", "country_names"))
    car = pm.Beta("car",  alpha=car_alpha, beta=car_beta, dims= ("media_names","country_names"))

    
    media_adstock = pm.Deterministic(
            name=f"media_adstock", 
            var=geometric_adstock(x=media, alpha=car[country], l_max=12), 
            dims=("obs_id", "media_names"))
    media_adstock_saturated = pm.Deterministic(
            name=f"media_adstock_saturated",
            var=logistic_saturation(x=media_adstock, lam = sat[country]),
            dims=("obs_id", "media_names"))
    regression_media = pm.Deterministic(
            name=f"regression_media", 
            var= media_adstock_saturated * coef[country], 
            dims = ('obs_id', 'media_names'))

    base = pm.Exponential(f"base", lam=base_lam, dims = "country_names")
    noise = pm.Exponential(f"noise", lam=0.001, dims = "country_names")

    sales = pm.Normal(
            "sales",
            mu=regression_media.sum(axis=-1) + base[country],
            sigma=noise[country],
            observed=y.values,
        )

I am sure the issue lies in my struggling to understand exactly how tensors work in this context so any help would be greatly appreciated.