code: def impute_missing_values(
      model,
      observed_time_series,
      parameter_samples,
      include_observation_noise=False,
      timesteps_are_event_shape=True,
  ):
      """Runs posterior inference to impute the missing values in a time series.
  
      This method computes the posterior marginals `p(latent state | observations)`,
      given the time series at observed timesteps (a missingness mask should
      be specified using `tfp.sts.MaskedTimeSeries`). It pushes this posterior back
      through the observation model to impute a predictive distribution on the
      observed time series. At unobserved steps, this is an imputed value; at other
      steps it is interpreted as the model's estimate of the underlying noise-free
      series.
  
      Args:
        model: `tfp.sts.Sum` instance defining an additive STS model.
        observed_time_series: `float` `Tensor` of shape
          `concat([sample_shape, model.batch_shape, [num_timesteps, 1]])` where
          `sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
          dimension may (optionally) be omitted if `num_timesteps > 1`. Any `NaN`s
            are interpreted as missing observations; missingness may be also be
            explicitly specified by passing a `tfp.sts.MaskedTimeSeries` instance.
        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.
        include_observation_noise: If `False`, the imputed uncertainties
          represent the model's estimate of the noise-free time series at each
          timestep. If `True`, they represent the model's estimate of the range of
          values that could be *observed* at each timestep, including any i.i.d.
          observation noise.
          Default value: `False`.
        timesteps_are_event_shape: Deprecated, for backwards compatibility only.
          If `False`, the predictive distribution will return per-timestep
          probabilities
          Default value: `True`.
  
      Returns:
        imputed_series_dist: a `tfd.MixtureSameFamily` instance with event shape
          `[num_timesteps] if timesteps_are_event_shape else []` and
          batch shape `concat([sample_shape, model.batch_shape,
          [] if timesteps_are_event_shape else [num_timesteps])`, with
          `num_posterior_draws` mixture components.
  
      #### Example
  
      To specify a time series with missing values, use `tfp.sts.MaskedTimeSeries`:
  
      ```python
      time_series_with_nans = [-1., 1., np.nan, 2.4, np.nan, 5]
      observed_time_series = tfp.sts.MaskedTimeSeries(
        time_series=time_series_with_nans,
        is_missing=tf.math.is_nan(time_series_with_nans))
      ```
  
      Masked time series can be passed to `tfp.sts` methods in place of a
      `observed_time_series` `Tensor`:
  
      ```python
      # Build model using observed time series to set heuristic priors.
      linear_trend_model = tfp.sts.LocalLinearTrend(
        observed_time_series=observed_time_series)
      model = tfp.sts.Sum([linear_trend_model],
                          observed_time_series=observed_time_series)
  
      # Fit model to data
      parameter_samples, _ = tfp.sts.fit_with_hmc(model, observed_time_series)
      ```
  
      After fitting a model, `impute_missing_values` will return a distribution
      ```python
      # Impute missing values
      imputed_series_distribution = tfp.sts.impute_missing_values(
        model, observed_time_series, parameter_samples=parameter_samples)
      print('imputed means and stddevs: ',
            imputed_series_distribution.mean(),
            imputed_series_distribution.stddev())
      ```
  
      """
      with tf.name_scope("impute_missing_values"):
          [
              observed_time_series,
              mask,
          ] = sts_util.canonicalize_observed_time_series_with_mask(
              observed_time_series
          )
  
          # Run smoothing over the training timesteps to extract the
          # predictive means and variances.
          num_timesteps = ps.dimension_size(observed_time_series, -2)
          lgssm = tfe_util.JitPublicMethods(
              model.make_state_space_model(
                  num_timesteps=num_timesteps,
                  param_vals=parameter_samples,
              ),
              trace_only=True,
          )  # Avoid eager overhead w/o introducing XLA dependence.
          posterior_means, posterior_covs = lgssm.posterior_marginals(
              observed_time_series, mask=mask
          )
  
          (
              observation_means,
              observation_covs,
          ) = lgssm.latents_to_observations(
              latent_means=posterior_means, latent_covs=posterior_covs
          )
  
          if not include_observation_noise:
              # Extract just the variance of observation noise by pushing forward
              # zero-variance latents.
              _, observation_noise_covs = lgssm.latents_to_observations(
                  latent_means=posterior_means,
                  latent_covs=tf.zeros_like(posterior_covs),
              )
              # Subtract out the observation noise that was added in the original
              # pushforward. Note that this could cause numerical issues if the
              # observation noise is very large. If this becomes an issue we could
              # avoid the subtraction by plumbing `include_observation_noise` through
              # `lgssm.latents_to_observations`.
              observation_covs -= observation_noise_covs
  
          # Squeeze dims to convert from LGSSM's event shape `[num_timesteps, 1]`
          # to a scalar time series.
          imputed_values_dist = sts_util.mix_over_posterior_draws(
              means=observation_means[..., 0],
              variances=observation_covs[..., 0, 0],
          )
          if timesteps_are_event_shape:
              imputed_values_dist = independent.Independent(
                  imputed_values_dist, reinterpreted_batch_ndims=1
              )
          return imputed_values_dist
  
file_path: probability/tensorflow_probability/python/sts/forecast.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import mixture_same_family
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.experimental import util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.sts.internal import util
  from tensorflow.python.util import deprecation
label_desc: Use image node_id 4 for calling a global function with example usage: impute_missing_values(model, observed_time_series, parameter_samples, include_observation_noise, timesteps_are_event_shape) and returns: imputed_values_dist