code: def test_one_step_predictive_correctness(self):
      observed_time_series_ = np.array(
          [1.0, -1.0, -3.0, 4.0, 0.5, 2.0, 1.0, 3.0]
      )
      observed_time_series = self._build_tensor(observed_time_series_)
      model = self._build_model(
          observed_time_series, constant_offset=0.0
      )  # Simplifies analytic calculations.
  
      drift_scale = 0.1
      observation_noise_scale = 0.01
      params = {
          "seasonal_drift_scale"
          if JAX_MODE
          else "seasonal/_drift_scale": self._build_tensor(
              [drift_scale]
          ),
          "observation_noise_scale": self._build_tensor(
              [observation_noise_scale]
          ),
      }
  
      @tf.function(autograph=False, jit_compile=tf.executing_eagerly())
      def _run():
          onestep_dist = forecast.one_step_predictive(
              model,
              observed_time_series,
              timesteps_are_event_shape=False,
              parameter_samples=params,
          )
          return onestep_dist.mean(), onestep_dist.stddev()
  
      onestep_mean, onestep_scale = _run()
  
      # Since Seasonal is just a set of interleaved random walks, it's
      # straightforward to compute the forecast analytically.
      # For the first (num_seasons - 1) steps, the one-step-ahead
      # forecast mean/scale are just the prior Normal(0., 1.). After that,
      # the predicted `n`th step depends on the posterior from the
      # `n - num_seasons` step.
      num_seasons = 4
      effect_posterior_precision = (
          1.0 + 1 / observation_noise_scale**2
      )
      effect_posterior_means = (
          observed_time_series_[:num_seasons]
          / observation_noise_scale**2
      ) / effect_posterior_precision
      effect_posterior_variance = 1 / effect_posterior_precision
      observation_predictive_variance = (
          effect_posterior_variance
          + drift_scale**2
          + observation_noise_scale**2
      )
  
      expected_onestep_mean = np.concatenate(
          [np.zeros(4), effect_posterior_means]
      )
      expected_onestep_scale = np.concatenate(
          [
              [np.sqrt(1.0**2 + observation_noise_scale**2)] * 4,
              [np.sqrt(observation_predictive_variance)] * 4,
          ]
      )
      self.assertAllClose(onestep_mean, expected_onestep_mean)
      self.assertAllClose(onestep_scale, expected_onestep_scale)
  
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 2 for calling the _ForecastTest obj's underlying member method code with example usage: obj.test_one_step_predictive_correctness() and returns: onestep_dist,  onestep_dist