code: def test_impute_missing(self):
      num_timesteps = 7
      time_series_with_nans = self._build_tensor(
          [-1.0, 1.0, np.nan, 2.4, np.nan, np.nan, 2.0]
      )
      observed_time_series = missing_values_util.MaskedTimeSeries(
          time_series=time_series_with_nans,
          is_missing=tf.math.is_nan(time_series_with_nans),
      )
  
      # Build model with a near-uniform prior on the initial effect. In principle
      # we should use a large scale like 1e8 here, but we use 1e2 because
      # increasing the scale triggers numerical issues with Kalman smoothing
      # described in b/138414045.
      model = self._build_model(
          observed_time_series, initial_effect_prior_scale=1e2
      )
  
      # Impute values using manually-set parameters, which will allow us to
      # compute the expected results analytically.
      drift_scale = 1.0
      noise_scale = 0.1
      parameter_samples = {
          "seasonal_drift_scale"
          if JAX_MODE
          else "seasonal/_drift_scale": self._build_tensor(
              [drift_scale]
          ),
          "observation_noise_scale": self._build_tensor([noise_scale]),
      }
  
      @tf.function(autograph=False, jit_compile=tf.executing_eagerly())
      def _run():
          imputed_series_dist = forecast.impute_missing_values(
              model,
              observed_time_series,
              parameter_samples,
              timesteps_are_event_shape=False,
          )
          imputed_noisy_series_dist = forecast.impute_missing_values(
              model,
              observed_time_series,
              parameter_samples,
              timesteps_are_event_shape=False,
              include_observation_noise=True,
          )
          mean, stddev = (
              imputed_series_dist.mean(),
              imputed_series_dist.stddev(),
          )
          noisy_mean, noisy_stddev = [
              imputed_noisy_series_dist.mean(),
              imputed_noisy_series_dist.stddev(),
          ]
          return (
              imputed_noisy_series_dist,
              mean,
              stddev,
              noisy_mean,
              noisy_stddev,
          )
  
      (
          imputed_noisy_series_dist,
          mean,
          stddev,
          noisy_mean,
          noisy_stddev,
      ) = _run()
  
      self.assertAllEqual(
          imputed_noisy_series_dist.batch_shape_tensor(),
          [num_timesteps],
      )
  
      # Compare imputed mean to expected mean.
      self.assertAllClose(
          mean, [-1.0, 1.0, 2.0, 2.4, -1.0, 1.0, 2.0], atol=1e-2
      )
      self.assertAllClose(mean, noisy_mean, atol=1e-2)
  
      # Compare imputed stddevs to expected stddevs.
      drift_plus_noise_scale = np.sqrt(
          noise_scale**2 + drift_scale**2
      )
      expected_stddev = np.array(
          [
              noise_scale,
              noise_scale,
              drift_plus_noise_scale,
              noise_scale,
              drift_plus_noise_scale,
              drift_plus_noise_scale,
              noise_scale,
          ]
      )
      self.assertAllClose(stddev, expected_stddev, atol=1e-2)
      self.assertAllClose(
          noisy_stddev,
          tf.sqrt(stddev**2 + noise_scale**2),
          atol=1e-2,
      )
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.sts import fitting
  from tensorflow_probability.python.sts import forecast
  from tensorflow_probability.python.sts.components import seasonal
  from tensorflow_probability.python.sts.components import sum
  from tensorflow_probability.python.sts.internal import missing_values_util
label_desc: Use image node_id 8 for calling the _ForecastTest obj's underlying member method code with example usage: obj.test_impute_missing() and returns: imputed_noisy_series_dist,  mean,  stddev,  noisy_mean,  noisy_stddev