code: def one_step_predictive(
      model,
      posterior_samples,
      num_forecast_steps=0,
      original_mean=0.0,
      original_scale=1.0,
      thin_every=1,
      use_zero_step_prediction=False,
  ):
      """Constructs a one-step-ahead predictive distribution at every timestep.
  
      Unlike the generic `tfp.sts.one_step_predictive`, this method uses the
      latent levels from Gibbs sampling to efficiently construct a predictive
      distribution that mixes over posterior samples. The predictive distribution
      may also include additional forecast steps.
  
      This method returns the predictive distributions for each timestep given
      previous timesteps and sampled model parameters, `p(observed_time_series[t] |
      observed_time_series[:t], weights, observation_noise_scale)`. Note that the
      posterior values of the weights and noise scale will in general be informed
      by observations from all timesteps *including the step being predicted*, so
      this is not a strictly kosher probabilistic quantity, but in general we assume
      that it's close, i.e., that the step being predicted had very small individual
      impact on the overall parameter posterior.
  
      Args:
        model: A `tfd.sts.StructuralTimeSeries` model instance. This must be of the
          form constructed by `build_model_for_gibbs_sampling`.
        posterior_samples: A `GibbsSamplerState` instance in which each element is a
          `Tensor` with initial dimension of size `num_samples`.
        num_forecast_steps: Python `int` number of additional forecast steps to
          append.
          Default value: `0`.
        original_mean: Optional scalar float `Tensor`, added to the predictive
          distribution to undo the effect of input normalization.
          Default value: `0.`
        original_scale: Optional scalar float `Tensor`, used to rescale the
          predictive distribution to undo the effect of input normalization.
          Default value: `1.`
        thin_every: Optional Python `int` factor by which to thin the posterior
          samples, to reduce complexity of the predictive distribution. For example,
          if `thin_every=10`, every `10`th sample will be used.
          Default value: `1`.
        use_zero_step_prediction: If true, instead of using the local level
          and trend from the timestep before, just use the local level from the
          same timestep.
  
      Returns:
        predictive_dist: A `tfd.MixtureSameFamily` instance of event shape
          `[num_timesteps + num_forecast_steps]` representing the predictive
          distribution of each timestep given previous timesteps.
      """
      dtype = dtype_util.common_dtype(
          [
              posterior_samples.level_scale,
              posterior_samples.observation_noise_scale,
              posterior_samples.level,
              original_mean,
              original_scale,
          ],
          dtype_hint=tf.float32,
      )
      num_observed_steps = prefer_static.shape(posterior_samples.level)[
          -1
      ]
  
      original_mean = tf.convert_to_tensor(original_mean, dtype=dtype)
      original_scale = tf.convert_to_tensor(original_scale, dtype=dtype)
      thinned_samples = tf.nest.map_structure(
          lambda x: x[::thin_every], posterior_samples
      )
  
      if (
          prefer_static.rank_from_shape(  # If no slope was inferred, treat as zero.
              prefer_static.shape(thinned_samples.slope)
          )
          <= 1
      ):
          thinned_samples = thinned_samples._replace(
              slope=tf.zeros_like(thinned_samples.level),
              slope_scale=tf.zeros_like(thinned_samples.level_scale),
          )
  
      num_steps_from_last_observation = tf.concat(
          [
              (
                  tf.zeros([num_observed_steps], dtype=dtype)
                  if use_zero_step_prediction
                  else tf.ones([num_observed_steps], dtype=dtype)
              ),
              tf.range(1, num_forecast_steps + 1, dtype=dtype),
          ],
          axis=0,
      )
  
      # The local linear trend model expects that the level at step t + 1 is equal
      # to the level at step t, plus the slope at time t - 1,
      # plus transition noise of scale 'level_scale' (which we account for below).
      if num_forecast_steps > 0:
          num_batch_dims = (
              prefer_static.rank_from_shape(
                  prefer_static.shape(thinned_samples.level)
              )
              - 2
          )
          # All else equal, the current level will remain stationary.
          forecast_level = tf.tile(
              thinned_samples.level[..., -1:],
              tf.concat(
                  [
                      tf.ones([num_batch_dims + 1], dtype=tf.int32),
                      [num_forecast_steps],
                  ],
                  axis=0,
              ),
          )
          # If the model includes slope, the level will steadily increase.
          forecast_level += thinned_samples.slope[..., -1:] * tf.range(
              1.0, num_forecast_steps + 1.0, dtype=forecast_level.dtype
          )
  
      level_pred = tf.concat(
          (
              [thinned_samples.level]
              if use_zero_step_prediction
              else [  # pylint:disable=g-long-ternary
                  thinned_samples.level[..., :1],  # t == 0
                  (
                      thinned_samples.level[..., :-1]
                      + thinned_samples.slope[..., :-1]
                  )  # 1 <= t < T. Constructs the next level from previous level
                  # and previous slope.
              ]
          )
          + ([forecast_level] if num_forecast_steps > 0 else []),
          axis=-1,
      )
  
      # Calculate the seasonal effect at each time point.
      (
          _,
          _,
          seasonal_indices_and_components,
      ) = _get_components_from_model(model)
      if seasonal_indices_and_components and num_forecast_steps > 0:
          raise NotImplementedError(
              "Forecasting with one_step_predictive is not supported with "
              "seasonality. Instead, use fit_with_gibbs_sampling with missing "
              "values."
          )
      seasonal_pred = _compute_seasonal_effect_from_levels(
          thinned_samples.seasonal_levels,
          [x[1] for x in seasonal_indices_and_components],
          default_dtype=dtype,
      )
  
      design_matrix = _get_design_matrix(model)
      if design_matrix is not None:
          design_matrix = design_matrix.to_dense()[
              : num_observed_steps + num_forecast_steps
          ]
          regression_effect = tf.linalg.matvec(
              design_matrix, thinned_samples.weights
          )
      else:
          regression_effect = 0
  
      y_mean = (
          level_pred + seasonal_pred + regression_effect
      ) * original_scale[..., tf.newaxis] + original_mean[
          ..., tf.newaxis
      ]
  
      # To derive a forecast variance, including slope uncertainty, let
      #  `r[:k]` be iid Gaussian RVs with variance `level_scale**2` and `s[:k]` be
      # iid Gaussian RVs with variance `slope_scale**2`. Then the forecast level at
      # step `T + k` can be written as
      #   (level[T] +           # Last known level.
      #    r[0] + ... + r[k] +  # Sum of random walk terms on level.
      #    slope[T] * k         # Contribution from last known slope.
      #    (k - 1) * s[0] +     # Contributions from random walk terms on slope.
      #    (k - 2) * s[1] +
      #    ... +
      #    1 * s[k - 1])
      # which has variance of
      #  (level_scale**2 * k +
      #   slope_scale**2 * ( (k - 1)**2 +
      #                      (k - 2)**2 +
      #                      ... + 1 ))
      # Here the `slope_scale` coefficient is the `k - 1`th square pyramidal
      # number [1], which is given by
      #  (k - 1) * k * (2 * k - 1) / 6.
      #
      # [1] https://en.wikipedia.org/wiki/Square_pyramidal_number
      variance_from_level = (
          thinned_samples.level_scale[..., tf.newaxis] ** 2
          * num_steps_from_last_observation
      )
      variance_from_slope = (
          thinned_samples.slope_scale[..., tf.newaxis] ** 2
          * (
              (num_steps_from_last_observation - 1)
              * num_steps_from_last_observation
              * (2 * num_steps_from_last_observation - 1)
          )
          / 6.0
      )
      y_scale = original_scale * tf.sqrt(
          thinned_samples.observation_noise_scale[..., tf.newaxis] ** 2
          + variance_from_level
          + variance_from_slope
      )
  
      num_posterior_draws = prefer_static.shape(y_mean)[0]
      return mixture_same_family.MixtureSameFamily(
          mixture_distribution=categorical.Categorical(
              logits=tf.zeros([num_posterior_draws], dtype=y_mean.dtype)
          ),
          components_distribution=normal.Normal(
              loc=dist_util.move_dimension(y_mean, 0, -1),
              scale=dist_util.move_dimension(y_scale, 0, -1),
          ),
      )
  
imports: import collections
  import itertools
  import numpy
  import six
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.bijectors import square
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import gamma
  from tensorflow_probability.python.distributions import inverse_gamma
  from tensorflow_probability.python.distributions import mixture_same_family
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_linear_operator
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import normal_conjugate_posteriors
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.experimental.distributions import mvn_precision_factor_linop
  from tensorflow_probability.python.experimental.sts_gibbs import dynamic_spike_and_slab
  from tensorflow_probability.python.experimental.sts_gibbs import sample_parameters
  from tensorflow_probability.python.experimental.sts_gibbs import spike_and_slab
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.sts.components import local_level
  from tensorflow_probability.python.sts.components import local_linear_trend
  from tensorflow_probability.python.sts.components import regression
  from tensorflow_probability.python.sts.components import seasonal
  from tensorflow_probability.python.sts.components import sum
  from tensorflow_probability.python.sts.internal import util
label_desc: Use image node_id 8 for calling a global function with example usage: one_step_predictive(model, posterior_samples, num_forecast_steps, original_mean, original_scale, thin_every, use_zero_step_prediction) and returns: mixture_same_family