tfp.sts.decompose_by_component fails when time series batch_size > 1

I am trying to train a structural time series model on several time series at once (batch_size = 3 in this example). If I understand correctly, this should be doable by constructing a time series matrix of shape (batch_size, num_timesteps).

Here is the model definition:

def build_model(observed_data, feature):
   monthly_effect = sts.Seasonal(
        num_seasons=12,
        num_steps_per_season=1,
        drift_scale_prior=tfd.Normal(loc=0, scale=0.01),
        constrain_mean_effect_to_zero=False,
        observed_time_series=observed_data,
        name='monthly_effect'
    )

    regressor_effect = sts.LinearRegression(
        design_matrix=feature - np.mean(feature),
        name='regressor_effect'
    )
    
    model = sts.Sum([monthly_effect,
                     regressor_effect
                    ],
                    observed_time_series=observed_data
                    )
    return model

The feature regressor has a shape of (num_timesteps, batch_size), which is the transpose of the shape of the time series matrix. It struck me as odd that that was the case, but the optimization (see below) only succeeded with those input dimensions.

I can then train the model without obvious issues using

def train_model(model, observed_data, variational_posteriors, num_variational_steps = 400):
    
    optimizer = tf.optimizers.Adam(learning_rate=.1)
    @tf.function(experimental_compile=True)
    def train():
        elbo_loss_curve = tfp.vi.fit_surrogate_posterior(
            target_log_prob_fn=model.joint_log_prob(
                observed_time_series=observed_data),
        surrogate_posterior=variational_posteriors,
        optimizer=optimizer,
        num_steps=num_variational_steps)
        return elbo_loss_curve

    elbo_loss_curve = train()
    return elbo_loss_curve, variational_posteriors

I then get an elbo curve for each constituent time series, and they all look good.

Sampling from the posterior also seems to work:

posterior_samples = variational_posteriors.sample(100)

Which has samples of shape (100, batch_size) for each parameter, as expected.

However, if I then try to decompose by component, I run into an error:

component_dists = sts.decompose_by_component(
    model,
    observed_time_series=observed_data,
    parameter_samples=posterior_samples)

results in:

InvalidArgumentError: Incompatible shapes: [100,1] vs. [3,1] [Op:AddV2] name: decompose_by_component/smooth/forward_filter/add/

The "3" in the error refers to the batch_size (which was 3 in this case). I cannot figure out how to reshape my inputs to make the decomposition work.
Any ideas?

Start at sts.decompose_by_component with the batched observed_time_series, model, and posterior_samples shown in the report. Trace the shape handling into the smooth/forward_filter operation named in the error, then reproduce the batch_size=3 case. Done means decomposition accepts these batched inputs and returns component distributions without the incompatible-shapes error.