code: def test_sampled_latents_have_correct_marginals(
      self,
      use_slope,
      tiled_residual_offset=None,
      seasonal_components_and_drift_scales_factory=None,
  ):
      if seasonal_components_and_drift_scales_factory is None:
          seasonal_components_and_drift_scales = []
      else:
          # Use a factory since distributions in JAX can not be initialized
          # before main runs.
          seasonal_components_and_drift_scales = (
              seasonal_components_and_drift_scales_factory()
          )
  
      seed = test_util.test_seed(sampler_type="stateless")
      residuals_seed, is_missing_seed, level_seed = samplers.split_seed(
          seed, 3, "test_sampled_level_has_correct_marginals"
      )
  
      # Have enough timesteps that seasonal effects can be observed.
      num_timesteps = 100
      batch_shape = (3, 1)
  
      residual_offset = 0.0
      if tiled_residual_offset:
          residual_offset = tf.convert_to_tensor(
              [
                  tiled_residual_offset[x % len(tiled_residual_offset)]
                  for x in range(num_timesteps)
              ]
          )
  
      observed_residuals = samplers.normal(
          batch_shape + (num_timesteps,), seed=residuals_seed
      ) + (residual_offset)
      is_missing = (
          samplers.uniform(
              batch_shape + (num_timesteps,), seed=is_missing_seed
          )
          > 0.8
      )
      level_scale = 1.5 * tf.ones(batch_shape)
      observation_noise_scale = 0.2 * tf.ones(batch_shape)
  
      if use_slope:
          initial_state_prior = mvn_diag.MultivariateNormalDiag(
              loc=[-30.0, 2.0], scale_diag=[1.0, 0.2]
          )
          slope_scale = 0.5 * tf.ones(batch_shape)
          level_ssm = (
              local_linear_trend.LocalLinearTrendStateSpaceModel(
                  num_timesteps=num_timesteps,
                  initial_state_prior=initial_state_prior,
                  observation_noise_scale=observation_noise_scale,
                  level_scale=level_scale,
                  slope_scale=slope_scale,
              )
          )
      else:
          initial_state_prior = mvn_diag.MultivariateNormalDiag(
              loc=[-30.0], scale_diag=[100.0]
          )
          slope_scale = None
          level_ssm = local_level.LocalLevelStateSpaceModel(
              num_timesteps=num_timesteps,
              initial_state_prior=initial_state_prior,
              observation_noise_scale=observation_noise_scale,
              level_scale=level_scale,
          )
  
      seasonal_ssms = []
      for (
          seasonal_component,
          drift_scale,
      ) in seasonal_components_and_drift_scales:
          seasonal_ssms.append(
              seasonal_component.make_state_space_model(
                  num_timesteps=num_timesteps,
                  initial_state_prior=seasonal_component.initial_state_prior,
                  param_vals={},
                  drift_scale=drift_scale,
                  # Set this to exactly 0, since we want the observation noise
                  # of the entire additive space to match observation_noise_scale,
                  # but that is already specified on the local ssm.
                  observation_noise_scale=0.0,
              )
          )
  
      ssm = sum_lib.AdditiveStateSpaceModel([level_ssm] + seasonal_ssms)
  
      posterior_means, posterior_covs = ssm.posterior_marginals(
          observed_residuals[..., tf.newaxis], mask=is_missing
      )
      latents_samples = gibbs_sampler._resample_latents(
          observed_residuals=observed_residuals,
          level_scale=level_scale,
          slope_scale=slope_scale,
          observation_noise_scale=observation_noise_scale,
          initial_state_prior=initial_state_prior,
          is_missing=is_missing,
          sample_shape=10000,
          seed=level_seed,
          seasonal_components=[
              x[0] for x in seasonal_components_and_drift_scales
          ],
          seasonal_drift_scales=tf.convert_to_tensor(
              value=[
                  x[1] for x in seasonal_components_and_drift_scales
              ],
              name="drift_scale",
          ),
      )
  
      (
          posterior_means_,
          posterior_covs_,
          latents_means_,
          latents_covs_,
      ) = self.evaluate(
          (
              posterior_means,
              posterior_covs,
              tf.reduce_mean(latents_samples, axis=0),
              sample_stats.covariance(
                  latents_samples, sample_axis=0, event_axis=-1
              ),
          )
      )
      # TODO(axch, cgs): Can we use assertAllMeansClose here?  The
      # latents_samples are presumably not IID across axis=0, so the
      # statistical assumptions are not satisfied.
      self.assertAllClose(latents_means_, posterior_means_, atol=0.1)
      self.assertAllClose(latents_covs_, posterior_covs_, atol=0.1)
  
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 9 for calling the GibbsSamplerTests obj's underlying member method code with example usage: obj.test_sampled_latents_have_correct_marginals(use_slope, tiled_residual_offset, seasonal_components_and_drift_scales_factory) without return types