code: def _build_sampler_loop_body(
      model,
      observed_time_series,
      is_missing=None,
      default_pseudo_observations=None,
      experimental_use_dynamic_cholesky=False,
      experimental_use_weight_adjustment=False,
  ):
      """Builds a Gibbs sampler for the given model and observed data.
  
      Args:
        model: A `tf.sts.StructuralTimeSeries` model instance. This must be of the
          form constructed by `build_model_for_gibbs_sampling`.
        observed_time_series: Float `Tensor` time series of shape `[...,
          num_timesteps]`.
        is_missing: Optional `bool` `Tensor` of shape `[..., num_timesteps]`. A
          `True` value indicates that the observation for that timestep is missing.
        default_pseudo_observations: Optional scalar float `Tensor` Controls the
          number of pseudo-observations for the prior precision matrix over the
          weights.
        experimental_use_dynamic_cholesky: Optional bool - in case of spike and slab
          sampling, will dynamically select the subset of the design matrix with
          active features to perform the Cholesky decomposition. This may provide
          a speedup when the number of true features is small compared to the size
          of the design matrix.
        experimental_use_weight_adjustment: Optional bool - use a nonstandard
          update for the posterior precision of the weight in case of a spike and
          slab sampler.
  
      Returns:
        sampler_loop_body: Python callable that performs a single cycle of Gibbs
          sampling. Its first argument is a `GibbsSamplerState`, and it returns a
          new `GibbsSamplerState`. The second argument (passed by `tf.scan`) is
          ignored.
      """
      if JAX_MODE and experimental_use_dynamic_cholesky:
          raise ValueError(
              "Dynamic Cholesky decomposition not supported in JAX"
          )
      (
          (_, level_component),
          (_, regression_component),
          seasonal_indices_and_components,
      ) = _get_components_from_model(model)
      seasonal_components = [
          x[1] for x in seasonal_indices_and_components
      ]
      model_has_slope = isinstance(
          level_component, local_linear_trend.LocalLinearTrend
      )
  
      if regression_component is not None:
          model_has_spike_slab_regression = isinstance(
              regression_component, SpikeAndSlabSparseLinearRegression
          )
  
      if is_missing is not None:  # Ensure series does not contain NaNs.
          observed_time_series = tf.where(
              is_missing,
              tf.zeros_like(observed_time_series),
              observed_time_series,
          )
      num_observed_steps = prefer_static.shape(observed_time_series)[-1]
  
      design_matrix = _get_design_matrix(model)
      num_missing = 0.0
      if design_matrix is not None:
          design_matrix = design_matrix.to_dense()[:num_observed_steps]
          if is_missing is None:
              num_missing = 0.0
              is_missing = tf.zeros(num_observed_steps, dtype=bool)
          else:
              # Replace design matrix with zeros at unobserved timesteps. This ensures
              # they will not affect the posterior on weights.
              design_matrix = tf.where(
                  is_missing[..., tf.newaxis],
                  tf.zeros_like(design_matrix),
                  design_matrix,
              )
              num_missing = tf.reduce_sum(
                  tf.cast(is_missing, design_matrix.dtype), axis=-1
              )
  
      # Untransform scale priors -> variance priors by reaching thru Sqrt bijector.
      observation_noise_param = model.parameters[0]
      if "observation_noise" not in observation_noise_param.name:
          raise ValueError(
              "Model parameters {} do not match the expected sampler "
              "state.".format(model.parameters)
          )
      observation_noise_variance_prior = (
          observation_noise_param.prior.distribution
      )
      if model_has_slope:
          level_scale_variance_prior, slope_scale_variance_prior = [
              p.prior.distribution for p in level_component.parameters
          ]
      else:
          level_scale_variance_prior = level_component.parameters[
              0
          ].prior.distribution
  
      if regression_component:
          if model_has_spike_slab_regression:
              if experimental_use_dynamic_cholesky:
                  sampler = (
                      dynamic_spike_and_slab.DynamicSpikeSlabSampler
                  )
              else:
                  sampler = spike_and_slab.SpikeSlabSampler
              spike_and_slab_sampler = sampler(
                  design_matrix,
                  weights_prior_precision=regression_component._weights_prior_precision,  # pylint: disable=protected-access
                  nonzero_prior_prob=regression_component._sparse_weights_nonzero_prob,  # pylint: disable=protected-access
                  observation_noise_variance_prior_concentration=(
                      observation_noise_variance_prior.concentration
                  ),
                  observation_noise_variance_prior_scale=(
                      observation_noise_variance_prior.scale
                  ),
                  observation_noise_variance_upper_bound=(
                      # The given bound is for the scale, so it must be squared to get
                      # the upper bound for the variance.
                      tf.math.square(
                          observation_noise_variance_prior.upper_bound
                      )
                      if hasattr(
                          observation_noise_variance_prior,
                          "upper_bound",
                      )
                      else None
                  ),
                  num_missing=num_missing,
                  **(
                      {
                          "default_pseudo_observations": default_pseudo_observations
                      }
                      if default_pseudo_observations is not None
                      else {}
                  ),
              )
              # In case the nonzero probability is exactly one, any proposal with any
              # zero weights will have log prob of -infinity, so we will pin the
              # proposals to one.
              # TODO(colcarroll): Can we short-circuit the feature selection loop in
              # case this is `True`?
              pin_to_nonzero = tf.greater_equal(
                  regression_component._sparse_weights_nonzero_prob, 1.0
              )  # pylint: disable=protected-access
  
          else:
              weights_prior_scale = regression_component.parameters[
                  0
              ].prior.scale
  
      # Sub-selects in `forward_filter_sequential` take up a lot of the runtime
      # with a dynamic Cholesky, but compiling here seems to help.
      # TODO(b/234726324): Should this always be compiled?
      if experimental_use_dynamic_cholesky:
          resample_latents = tf.function(
              jit_compile=True, autograph=False
          )(_resample_latents)
          resample_scale = tf.function(
              jit_compile=True, autograph=False
          )(_resample_scale)
          resample_seasonal_scales = tf.function(
              jit_compile=True, autograph=False
          )(_resample_seasonal_scales)
      else:
          resample_latents = _resample_latents
          resample_scale = _resample_scale
          resample_seasonal_scales = _resample_seasonal_scales
  
      def sampler_loop_body(previous_sample, _):
          """Runs one sampler iteration, resampling all model variables."""
  
          (
              weights_seed,
              latent_seed,
              observation_noise_scale_seed,
              level_scale_seed,
              loop_seed,
              seasonal_seed,
          ) = samplers.split_seed(
              previous_sample.seed, n=6, salt="sampler_loop_body"
          )
          # Preserve backward-compatible seed behavior by splitting slope separately.
          (slope_scale_seed,) = samplers.split_seed(
              previous_sample.seed, n=1, salt="sampler_loop_body_slope"
          )
  
          if regression_component:
              # We encourage a reasonable initialization by sampling the weights first,
              # so at the first step they are regressed directly against the observed
              # time series. If we instead sampled the level first it might 'explain
              # away' some observed variation that we would ultimately prefer to explain
              # through the regression weights, because the level can represent
              # arbitrary variation, while the weights are limited to representing
              # variation in the subspace given by the design matrix.
              all_seasons_effect = _compute_seasonal_effect_from_levels(
                  previous_sample.seasonal_levels,
                  [x[1] for x in seasonal_indices_and_components],
                  default_dtype=previous_sample.level.dtype,
              )
              target_residuals = (
                  observed_time_series
                  - previous_sample.level
                  - all_seasons_effect
              )
              if model_has_spike_slab_regression:
                  if experimental_use_weight_adjustment:
                      previous_observation_noise_variance = tf.square(
                          previous_sample.observation_noise_scale
                      )
                  else:
                      previous_observation_noise_variance = 1.0
                  targets = tf.where(
                      is_missing,
                      tf.zeros_like(observed_time_series),
                      target_residuals,
                  )
                  (
                      observation_noise_variance,
                      weights,
                  ) = spike_and_slab_sampler.sample_noise_variance_and_weights(
                      initial_nonzeros=tf.math.logical_or(
                          tf.not_equal(previous_sample.weights, 0.0),
                          pin_to_nonzero,
                      ),
                      previous_observation_noise_variance=previous_observation_noise_variance,
                      targets=targets,
                      seed=weights_seed,
                  )
                  observation_noise_scale = tf.sqrt(
                      observation_noise_variance
                  )
  
              else:
                  weights = _resample_weights(
                      design_matrix=design_matrix,
                      target_residuals=target_residuals,
                      observation_noise_scale=previous_sample.observation_noise_scale,
                      weights_prior_scale=weights_prior_scale,
                      seed=weights_seed,
                  )
                  # Noise scale will be resampled below.
                  observation_noise_scale = (
                      previous_sample.observation_noise_scale
                  )
  
              regression_residuals = (
                  observed_time_series
                  - tf.linalg.matvec(design_matrix, weights)
              )
          else:
              # If there is no regression, then the entire timeseries is a residual.
              regression_residuals = observed_time_series
              # Noise scale will be resampled below.
              observation_noise_scale = (
                  previous_sample.observation_noise_scale
              )
              weights = previous_sample.weights
  
          latents = resample_latents(
              observed_residuals=regression_residuals,
              level_scale=previous_sample.level_scale,
              slope_scale=previous_sample.slope_scale
              if model_has_slope
              else None,
              observation_noise_scale=observation_noise_scale,
              initial_state_prior=level_component.initial_state_prior,
              is_missing=is_missing,
              seed=latent_seed,
              seasonal_components=seasonal_components,
              seasonal_drift_scales=previous_sample.seasonal_drift_scales,
          )
          level = latents[..., 0]
          level_residuals = level[..., 1:] - level[..., :-1]
          if model_has_slope:
              slope = latents[..., 1]
              level_residuals -= slope[..., :-1]
              slope_residuals = slope[..., 1:] - slope[..., :-1]
          next_latent_index = 2 if model_has_slope else 1
          if seasonal_components:
              num_seasonal_latents = sum(
                  [
                      component.latent_size
                      for component in seasonal_components
                  ]
              )
              # Seasonal levels are after the local latents.
              seasonal_levels = latents[
                  ...,
                  next_latent_index : next_latent_index
                  + num_seasonal_latents,
              ]
              next_latent_index += num_seasonal_latents
              seasonal_drift_scales = resample_seasonal_scales(
                  seasonal_components, seasonal_levels, seasonal_seed
              )
          else:
              # When there are no seasons, the lists will be empty and tf.stack
              # will fail. Thus, just re-use whatever (unchanging) shape was
              # initially set so it is unchanging on each iteration.
              seasonal_drift_scales = (
                  previous_sample.seasonal_drift_scales
              )
              seasonal_levels = previous_sample.seasonal_levels
          all_seasons_effect = _compute_seasonal_effect_from_levels(
              seasonal_levels,
              [x[1] for x in seasonal_indices_and_components],
              default_dtype=level.dtype,
          )
  
          number_of_latents = latents.shape[-1]
          if next_latent_index != number_of_latents:
              raise TypeError(
                  f"Some latent values were not used. {next_latent_index} were used, "
                  "but there were {number_of_latents} latents. Are you using "
                  "unsupported structual components?"
              )
  
          # Estimate level scale from the empirical changes in level.
          level_scale = resample_scale(
              prior=level_scale_variance_prior,
              observed_residuals=level_residuals,
              is_missing=None,
              seed=level_scale_seed,
          )
          if model_has_slope:
              slope_scale = resample_scale(
                  prior=slope_scale_variance_prior,
                  observed_residuals=slope_residuals,
                  is_missing=None,
                  seed=slope_scale_seed,
              )
          if not (
              regression_component and model_has_spike_slab_regression
          ):
              # Estimate noise scale from the residuals.
              observation_noise_scale = resample_scale(
                  prior=observation_noise_variance_prior,
                  observed_residuals=regression_residuals
                  - level
                  - all_seasons_effect,
                  is_missing=is_missing,
                  seed=observation_noise_scale_seed,
              )
  
          return GibbsSamplerState(
              observation_noise_scale=observation_noise_scale,
              level_scale=level_scale,
              slope_scale=(
                  slope_scale
                  if model_has_slope
                  else previous_sample.slope_scale
              ),
              weights=weights,
              level=level,
              slope=(
                  slope if model_has_slope else previous_sample.slope
              ),
              seasonal_drift_scales=seasonal_drift_scales,
              seasonal_levels=seasonal_levels,
              seed=loop_seed,
          )
  
      return sampler_loop_body
  
imports: import collections
  import itertools
  import numpy
  import six
  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 categorical
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import gamma
  from tensorflow_probability.python.distributions import inverse_gamma
  from tensorflow_probability.python.distributions import mixture_same_family
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_linear_operator
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import normal_conjugate_posteriors
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.experimental.distributions import mvn_precision_factor_linop
  from tensorflow_probability.python.experimental.sts_gibbs import dynamic_spike_and_slab
  from tensorflow_probability.python.experimental.sts_gibbs import sample_parameters
  from tensorflow_probability.python.experimental.sts_gibbs import spike_and_slab
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import samplers
  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 sum
  from tensorflow_probability.python.sts.internal import util
label_desc: Use image node_id 15 for calling a global function with example usage: _build_sampler_loop_body(model, observed_time_series, is_missing, default_pseudo_observations, experimental_use_dynamic_cholesky, experimental_use_weight_adjustment) and returns: sampler_loop_body, GibbsSamplerState