code: def _decompose_from_posterior_marginals(
      model,
      posterior_means,
      posterior_covs,
      parameter_samples,
      initial_step=0,
  ):
      """Utility method to decompose a joint posterior into components.
  
      Args:
        model: `tfp.sts.Sum` instance defining an additive STS model.
        posterior_means: float `Tensor` of shape `concat(
          [[num_posterior_draws], batch_shape, num_timesteps, latent_size])`
          representing the posterior mean over latents in an
          `AdditiveStateSpaceModel`.
        posterior_covs: float `Tensor` of shape `concat(
          [[num_posterior_draws], batch_shape, num_timesteps,
          latent_size, latent_size])`
          representing the posterior marginal covariances over latents in an
          `AdditiveStateSpaceModel`.
        parameter_samples: Python `list` of `Tensors` representing posterior
          samples of model parameters, with shapes `[concat([
          [num_posterior_draws], param.prior.batch_shape,
          param.prior.event_shape]) for param in model.parameters]`. This may
          optionally also be a map (Python `dict`) of parameter names to
          `Tensor` values.
        initial_step: optional `int` specifying the initial timestep of the
          decomposition.
  
      Returns:
        component_dists: A `collections.OrderedDict` instance mapping
          component StructuralTimeSeries instances (elements of `model.components`)
          to `tfd.Distribution` instances representing the posterior marginal
          distributions on the process modeled by each component. Each distribution
          has batch shape matching that of `posterior_means`/`posterior_covs`, and
          event shape of `[num_timesteps]`.
      """
  
      try:
          model.components
      except AttributeError:
          raise ValueError(
              "Model decomposed into components must be an instance of"
              "`tfp.sts.Sum` (passed model {})".format(model)
          )
  
      with tf.name_scope("decompose_from_posterior_marginals"):
          # Extract the component means/covs from the joint latent posterior.
          latent_sizes = [
              component.latent_size for component in model.components
          ]
          component_means = tf.split(
              posterior_means, latent_sizes, axis=-1
          )
          component_covs = _split_covariance_into_marginals(
              posterior_covs, latent_sizes
          )
  
          # Instantiate per-component state space models, and use them to push the
          # posterior means/covs through the observation model for each component.
          num_timesteps = dist_util.prefer_static_value(
              tf.shape(posterior_means)
          )[-2]
          component_ssms = model.make_component_state_space_models(
              num_timesteps=num_timesteps,
              param_vals=parameter_samples,
              initial_step=initial_step,
          )
          component_predictive_dists = collections.OrderedDict()
          for (
              component,
              component_ssm,
              component_mean,
              component_cov,
          ) in zip(
              model.components,
              component_ssms,
              component_means,
              component_covs,
          ):
              (
                  component_obs_mean,
                  component_obs_cov,
              ) = component_ssm.latents_to_observations(
                  latent_means=component_mean, latent_covs=component_cov
              )
  
              # Using the observation means and covs, build a mixture distribution
              # that integrates over the posterior draws.
              component_predictive_dists[
                  component
              ] = sts_util.mix_over_posterior_draws(
                  means=component_obs_mean[..., 0],
                  variances=component_obs_cov[..., 0, 0],
              )
      return component_predictive_dists
  
file_path: probability/tensorflow_probability/python/sts/decomposition.py
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental import util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.sts.internal import util
label_desc: Use image node_id 2 for calling a global function with example usage: _decompose_from_posterior_marginals(model, posterior_means, posterior_covs, parameter_samples, initial_step) and returns: component_predictive_dists