Error when creating coef_extra_features priors for multiple geos

[rodrigops123]

Hi, I'm having trouble when creating the coef_extra_features priors for 2 geos. I'm using a normal distribution as a probability function, however I'm attributing different loc and scale parameters depending on the extra feature I'm considering. This is a sample code:

extra_features_prior_scale = {"extra_feat_1": 0, "extra_feat_2": -1, "extra_feat_3": 1, ..., "extra_feat_12": 1}
extra_features_prior_loc = {"extra_feat_1": 1, "extra_feat_2": 1, "extra_feat_3": 1, ..., "extra_feat_12": 1}

extra_features_loc_priors = jnp.array(list(extra_features_prior_loc.values()) + [0])

extra_features_scale_priors = jnp.array(list(extra_features_prior_scale.values()) + [1])

custom_priors = {
    "half_max_effective_concentration": numpyro.distributions.Gamma(
        concentration=media_concentration_weights_priors,
        rate=media_concentration_weights_priors,
    ),
    "slope": numpyro.distributions.Gamma(
        concentration=slope_concentration_priors, rate=slope_rate_priors
    ),
    "lag_weight": numpyro.distributions.Beta(
        concentration1=lag_weight_high_priors, concentration0=lag_weight_low_priors
    ),
    "intercept": numpyro.distributions.Normal(
        loc=intercept_loc_priors, scale=intercept_scale_priors
    ),
    "coef_extra_features": numpyro.distributions.Normal(
        loc=extra_features_loc_priors,
        scale=extra_features_scale_priors,
   )
}

However, when training my algorithm, I'm getting the following error:

ValueError: Incompatible shapes for broadcasting: shapes=[(12, 2), (12, 12)]

I'm pretty sure this error comes from the coef_extra_features prior since I didn't have this until I decided to customize my priors on these features.