Multi-step predictions with STS

Problem

Hey, I'm working from the example STS notebook an I have question regarding making multi-step forecasts from a set of data. The notebook makes use of the tfp.sts.one_step_predictive, which returns "... the predictive distribution over observations at each time T, given observations up through time T-1"

For my problem, I need to make predictions 5 time-steps ahead, so I need the same functionality but instead generate predictive distribution over observations at each time T, given observations up through time T-5.

Question

Is there an efficient way to replicate the behaviour of tfp.sts.one_step_predictive, except have the output be based on observations up through T-n?

Attempted solution

In an attempt to implement it, I first do:

num_observed_steps = dist_util.prefer_static_value(
        tf.shape(observed_time_series))[-2]
observed_data_ssm = model.make_state_space_model(
        num_timesteps=num_observed_steps, param_vals=parameter_samples)
(_, _, _, predictive_means, predictive_covs, _, _
) = observed_data_ssm.forward_filter(observed_time_series, mask=mask)

and then use each pair of elements in predictive_means and predictive_covs as the initial_state_prior to forecast 5 days ahead, similarly to how it is implemented in the forecast function forecast.py. I then just collect the last forecast of each of these 5 day windows and return them as the predictions given a set.

This is extremely slow though, and I am still learning the theory so I'm not sure if this makes any sense or if it is the most efficient way to do it. Any pointers would be greatly appreciated!

Start with the STS example notebook and tensorflow_probability/python/sts/forecast.py, comparing one_step_predictive with the attempted forward_filter path. Trace how predictive means, covariances, and initial_state_prior are used for an n-step horizon; done means an agreed efficient approach or a clearly documented limitation.