Hierarchical Bayesian Unobserved Component Model

Question

I'm pretty new to TFP, and probabilistic programming in general, so apologies in advance.

I'm wanting to know how to build a Hierarchical UCM in TFP state space. The idea would be to do something like this (obviously this doesn't work)

# Specify a global regression parameter, all calculated regression parameters should fall around this.
global_loc = tfp.HalfNormal(100)

# Standard
trend_effect1 = tfp.sts.LocalLinearTrend(
    level_scale_prior=tfd.HalfNormal(100),
    slope_scale_prior=tfd.HalfNormal(100),
    observed_time_series=df_endo1.values,
    name='trend_effect'
)

# Have the prior weight centered around the global mean.
exo_effects1 = tfp.sts.LinearRegression(
    design_matrix=df_exog1.values,
    weights_prior=tfd.Normal(loc=global_loc, scale = 1),
    name='exo_effects',
)

# Standard
model1 = tfp.sts.Sum(
    [trend_effect1, exo_effects1],
    observed_time_series=df_endo1.values,
    observation_noise_scale_prior=tfp.distributions.HalfNormal(10)
)

# Standard
trend_effect2 = tfp.sts.LocalLinearTrend(
    level_scale_prior=tfd.HalfNormal(1),
    slope_scale_prior=tfd.Normal(loc = global_loc, scale = 1),
    observed_time_series=df_endo2.values,
    name='trend_effect'
)

# Same as above but this time for the second model with different endo and exog data.
exo_effects2 = tfp.sts.LinearRegression(
    design_matrix=df_exog2.values,
    weights_prior=tfd.Normal(loc=global_loc, scale = 1),
    name='exo_effects',
)

model2 = tfp.sts.Sum(
    [trend_effect2, exo_effects2],
    observed_time_series=df_endo2.values,
    observation_noise_scale_prior=tfp.distributions.HalfNormal(10)
)
...

So, we sample from some global mean value, and use it as the mean for the priors of regressors., this way the are similar but not the same.

This would be a pretty standard procedure if it weren't a UCM.

Happy to have suggestions outside of TFP if people think there are better options, or to use a Bayesian model based on the UCM from statsmodels (as I have seen this done before).

Thanks.

I've looked at TFPs Multilevel Model but am not sure how this would work for a UCM. I've also tried the code I wrote above but unsurprisingly it doesn't work.