code: def test_forecasts_match_reference(
      self,
      use_slope,
      num_chains,
      time_series_shift,
      use_zero_step_prediction=False,
      seasonal_components=None,
      use_num_forecast_steps_for_forecast=True,
      tiled_residual_offset=None,
      forecast_stddev_atol=None,
  ):
      if seasonal_components is None:
          seasonal_components = []
      dtype = tf.float32
      # pylint:disable=g-complex-comprehension
      seasonal_components = [
          seasonal.Seasonal(
              name=f"season{season_index}",
              num_seasons=num_seasons,
              drift_scale_prior=transformed_distribution.TransformedDistribution(
                  bijector=invert.Invert(square.Square()),
                  distribution=inverse_gamma.InverseGamma(
                      tf.constant(0.01, dtype=dtype),
                      tf.constant(0.0001, dtype=dtype),
                  ),
              ),
              initial_effect_prior=normal.Normal(
                  tf.constant(0.0, dtype=dtype),
                  tf.constant(1.0, dtype=dtype),
              ),
          )
          for season_index, num_seasons in enumerate(
              seasonal_components
          )
      ]
      # pylint:enable=g-complex-comprehension
      seed = test_util.test_seed()
      # Have enough steps to fit seasonality to.
      num_observed_steps = 100
      num_forecast_steps = 4
      num_results = 1000
  
      # Dividing the number of results with number of chains so we sample the same
      # total number of MCMC samples.
      if not tf.nest.is_nested(num_chains):
          num_results = num_results // num_chains
  
      batch_shape = [3]
      model, observed_time_series, is_missing = self._build_test_model(
          num_timesteps=num_observed_steps + num_forecast_steps,
          true_slope_scale=0.5 if use_slope else None,
          batch_shape=batch_shape,
          time_series_shift=time_series_shift,
          seasonal_components=seasonal_components,
          tiled_residual_offset=tiled_residual_offset,
      )
  
      if use_num_forecast_steps_for_forecast:
          time_series_to_fit = missing_values_util.MaskedTimeSeries(
              observed_time_series[
                  ..., :num_observed_steps, tf.newaxis
              ],
              is_missing[..., :num_observed_steps],
          )
      else:
          # Use missing values for the last N points instead of forecasting.
          time_series_to_fit = missing_values_util.MaskedTimeSeries(
              observed_time_series[..., tf.newaxis],
              tf.concat(
                  [
                      is_missing[..., :num_observed_steps],
                      tf.broadcast_to(
                          False,
                          is_missing.shape[:-1] + (num_forecast_steps,),
                      ),
                  ],
                  axis=-1,
              ),
          )
  
      @tf.function(autograph=False)
      def do_sampling():
          return gibbs_sampler.fit_with_gibbs_sampling(
              model,
              time_series_to_fit,
              num_chains=num_chains,
              num_results=num_results,
              num_warmup_steps=100,
              seed=seed,
          )
  
      samples = self.evaluate(do_sampling())
  
      def reshape_chain_and_sample(x):
          if np.ndim(x) > 2:
              return np.reshape(
                  x, [x.shape[0] * x.shape[1], *x.shape[2:]]
              )
          return x
  
      if not tf.nest.is_nested(num_chains):
          samples = tf.nest.map_structure(
              reshape_chain_and_sample, samples
          )
  
      predictive_dist = gibbs_sampler.one_step_predictive(
          model,
          samples,
          num_forecast_steps=(
              # If we are not using forecast steps for forecast, then the
              # predictions are already included through the masked inputs.
              num_forecast_steps
              if use_num_forecast_steps_for_forecast
              else 0
          ),
          thin_every=1,
          use_zero_step_prediction=use_zero_step_prediction,
      )
      predictive_mean, predictive_stddev = self.evaluate(
          (predictive_dist.mean(), predictive_dist.stddev())
      )
      self.assertAllEqual(
          predictive_mean.shape,
          batch_shape + [num_observed_steps + num_forecast_steps],
      )
      self.assertAllEqual(
          predictive_stddev.shape,
          batch_shape + [num_observed_steps + num_forecast_steps],
      )
  
      # big tolerance, but makes sure the predictive mean initializes near
      # the initial time series value
      self.assertAllClose(
          tf.reduce_mean(predictive_mean[:, 0]),
          observed_time_series[0, 0],
          atol=10.0,
      )
      parameter_samples = (
          gibbs_sampler.model_parameter_samples_from_gibbs_samples(
              model, samples
          )
      )
  
      # Note that although we expect the Gibbs-sampled forecasts to match a
      # reference implementation, we *don't* expect the one-step predictions to
      # match `tfp.sts.one_step_predictive`, because that makes predictions using
      # a filtered posterior (i.e., given only previous observations) whereas the
      # Gibbs-sampled latent `level`s will incorporate some information from
      # future observations.
      reference_forecast_dist = forecast(
          model,
          observed_time_series=observed_time_series[
              ..., :num_observed_steps
          ],
          parameter_samples=parameter_samples,
          num_steps_forecast=num_forecast_steps,
      )
  
      (
          reference_forecast_mean,
          reference_forecast_stddev,
      ) = self.evaluate(
          (
              reference_forecast_dist.mean()[..., 0],
              reference_forecast_dist.stddev()[..., 0],
          )
      )
  
      self.assertAllClose(
          predictive_mean[..., -num_forecast_steps:],
          reference_forecast_mean,
          atol=1.0 if use_slope else 0.3,
      )
      if forecast_stddev_atol is None:
          forecast_stddev_atol = 2.0 if use_slope else 1.00
      self.assertAllClose(
          predictive_stddev[..., -num_forecast_steps:],
          reference_forecast_stddev,
          atol=forecast_stddev_atol,
      )
  
imports: from absl.testing import parameterized
  import numpy
  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 inverse_gamma
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import student_t
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.distributions.linear_gaussian_ssm import linear_gaussian_update
  from tensorflow_probability.python.experimental.distributions import mvn_precision_factor_linop
  from tensorflow_probability.python.experimental.sts_gibbs import gibbs_sampler
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.stats import sample_stats
  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 semilocal_linear_trend
  from tensorflow_probability.python.sts.components import sum
  from tensorflow_probability.python.sts.forecast import forecast
  from tensorflow_probability.python.sts.internal import missing_values_util
  from tensorflow_probability.python.sts.internal import util
label_desc: Use image node_id 2 for calling the GibbsSamplerTests obj's underlying member method code with example usage: obj.test_forecasts_match_reference(use_slope, num_chains, time_series_shift, use_zero_step_prediction, seasonal_components, use_num_forecast_steps_for_forecast, tiled_residual_offset, forecast_stddev_atol) and returns: gibbs_sampler, x, np