code: def test_sampled_weights_follow_correct_distribution(self):
      seed = test_util.test_seed(sampler_type="stateless")
      (
          design_seed,
          true_weights_seed,
          sampled_weights_seed,
      ) = samplers.split_seed(
          seed, 3, "test_sampled_weights_follow_correct_distribution"
      )
      num_timesteps = 10
      num_features = 2
      batch_shape = [3, 1]
      design_matrix = self.evaluate(
          samplers.normal(
              batch_shape + [num_timesteps, num_features],
              seed=design_seed,
          )
      )
      true_weights = self.evaluate(
          samplers.normal(
              batch_shape + [num_features, 1], seed=true_weights_seed
          )
          * 10.0
      )
      targets = np.matmul(design_matrix, true_weights)
      is_missing = np.array(
          [
              False,
              False,
              False,
              True,
              True,
              False,
              False,
              True,
              False,
              False,
          ]
      )
      prior_scale = tf.convert_to_tensor(5.0)
      likelihood_scale = tf.convert_to_tensor(0.1)
  
      # Analytically compute the true posterior distribution on weights.
      valid_design_matrix = design_matrix[..., ~is_missing, :]
      valid_targets = targets[..., ~is_missing, :]
      num_valid_observations = tf.shape(valid_design_matrix)[-2]
      (
          weights_posterior_mean,
          weights_posterior_cov,
          _,
      ) = linear_gaussian_update(
          prior_mean=tf.zeros([num_features, 1]),
          prior_cov=tf.eye(num_features) * prior_scale**2,
          observation_matrix=tfl.LinearOperatorFullMatrix(
              valid_design_matrix
          ),
          observation_noise=mvn_diag.MultivariateNormalDiag(
              loc=tf.zeros([num_valid_observations]),
              scale_diag=likelihood_scale
              * tf.ones([num_valid_observations]),
          ),
          x_observed=valid_targets,
      )
  
      # Check that the empirical moments of sampled weights match the true values.
      sampled_weights = tf.vectorized_map(
          lambda seed: gibbs_sampler._resample_weights(  # pylint: disable=g-long-lambda
              design_matrix=tf.where(
                  is_missing[..., tf.newaxis],
                  tf.zeros_like(design_matrix),
                  design_matrix,
              ),
              target_residuals=targets[..., 0],
              observation_noise_scale=likelihood_scale,
              weights_prior_scale=tf.linalg.LinearOperatorScaledIdentity(
                  num_features, prior_scale
              ),
              seed=seed,
          ),
          samplers.split_seed(sampled_weights_seed, tf.constant(10000)),
      )
      sampled_weights_mean = tf.reduce_mean(sampled_weights, axis=0)
      centered_weights = (
          sampled_weights - weights_posterior_mean[..., 0]
      )
      sampled_weights_cov = tf.reduce_mean(
          centered_weights[..., :, tf.newaxis]
          * centered_weights[..., tf.newaxis, :],
          axis=0,
      )
  
      (
          sampled_weights_mean_,
          weights_posterior_mean_,
          sampled_weights_cov_,
          weights_posterior_cov_,
      ) = self.evaluate(
          (
              sampled_weights_mean,
              weights_posterior_mean[..., 0],
              sampled_weights_cov,
              weights_posterior_cov,
          )
      )
      self.assertAllClose(
          sampled_weights_mean_,
          weights_posterior_mean_,
          atol=0.01,
          rtol=0.05,
      )
      self.assertAllClose(
          sampled_weights_cov_,
          weights_posterior_cov_,
          atol=0.01,
          rtol=0.05,
      )
  
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 11 for calling the GibbsSamplerTests obj's underlying member method code with example usage: obj.test_sampled_weights_follow_correct_distribution() without return types