code: def fit_with_gibbs_sampling(
      model,
      observed_time_series,
      num_chains=(),
      num_results=2000,
      num_warmup_steps=200,
      initial_state=None,
      seed=None,
      default_pseudo_observations=None,
      experimental_use_dynamic_cholesky=False,
      experimental_use_weight_adjustment=False,
  ):
      """Fits parameters for an STS model using Gibbs sampling.
  
      Args:
        model: A `tfp.sts.StructuralTimeSeries` model instance return by
          `build_model_for_gibbs_fitting`.
        observed_time_series: `float` `Tensor` of shape [..., T, 1]` (omitting the
          trailing unit dimension is also supported when `T > 1`), specifying an
          observed time series. May optionally be an instance of
          `tfp.sts.MaskedTimeSeries`, which includes a mask `Tensor` to specify
          timesteps with missing observations.
        num_chains: Optional int to indicate the number of parallel MCMC chains.
          Default to an empty tuple to sample a single chain.
        num_results: Optional int to indicate number of MCMC samples.
        num_warmup_steps: Optional int to indicate number of MCMC samples.
        initial_state: A `GibbsSamplerState` structure of the initial states of the
          MCMC chains.
        seed: Optional `Python` `int` seed controlling the sampled values.
        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. *Note*: If this is true, neither batch shape nor
          `jit_compile` is supported.
        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:
        model: A `GibbsSamplerState` structure of posterior samples.
      """
      if not tf.nest.is_nested(num_chains):
          num_chains = [num_chains]
  
      [
          observed_time_series,
          is_missing,
      ] = sts_util.canonicalize_observed_time_series_with_mask(
          observed_time_series
      )
      dtype = observed_time_series.dtype
  
      # The canonicalized time series always has trailing dimension `1`,
      # because although LinearGaussianSSMs support vector observations, STS models
      # describe scalar time series only. For our purposes it'll be cleaner to
      # remove this dimension.
      observed_time_series = observed_time_series[..., 0]
      batch_shape = prefer_static.concat(
          [num_chains, prefer_static.shape(observed_time_series)[:-1]],
          axis=-1,
      )
      level_slope_shape = prefer_static.concat(
          [num_chains, prefer_static.shape(observed_time_series)],
          axis=-1,
      )
  
      # Treat a LocalLevel model as the special case of LocalLinearTrend where
      # the slope_scale is always zero.
      initial_slope_scale = 0.0
      initial_slope = 0.0
      (
          (_, level_component),
          _,
          seasonal_indices_and_components,
      ) = _get_components_from_model(model)
      if isinstance(
          level_component, local_linear_trend.LocalLinearTrend
      ):
          initial_slope_scale = 1.0 * tf.ones(batch_shape, dtype=dtype)
          initial_slope = tf.zeros(level_slope_shape, dtype=dtype)
  
      if initial_state is None:
          design_matrix = _get_design_matrix(model)
          weights = (
              tf.zeros(0, dtype=dtype)
              if design_matrix is None
              else tf.zeros(  # pylint:disable=g-long-ternary
                  prefer_static.concat(
                      [batch_shape, design_matrix.shape[-1:]], axis=0
                  ),
                  dtype=dtype,
              )
          )
          initial_state = GibbsSamplerState(
              observation_noise_scale=tf.ones(batch_shape, dtype=dtype),
              level_scale=tf.ones(batch_shape, dtype=dtype),
              slope_scale=initial_slope_scale,
              weights=weights,
              level=tf.zeros(level_slope_shape, dtype=dtype),
              slope=initial_slope,
              seed=None,  # Set below.
              seasonal_drift_scales=tf.ones(
                  prefer_static.concat(
                      [
                          batch_shape,
                          [len(seasonal_indices_and_components)],
                      ],
                      axis=0,
                  ),
                  dtype=dtype,
              ),
              seasonal_levels=tf.zeros(
                  get_seasonal_latents_shape(
                      observed_time_series, model, num_chains=num_chains
                  ),
                  dtype=dtype,
              ),
          )
  
      if isinstance(seed, six.integer_types):
          tf.random.set_seed(seed)
  
      # Always use the passed-in `seed` arg, ignoring any seed in the initial state.
      initial_state = initial_state._replace(
          seed=samplers.sanitize_seed(
              seed, salt="initial_GibbsSamplerState"
          )
      )
  
      sampler_loop_body = _build_sampler_loop_body(
          model=model,
          observed_time_series=observed_time_series,
          is_missing=is_missing,
          default_pseudo_observations=default_pseudo_observations,
          experimental_use_dynamic_cholesky=experimental_use_dynamic_cholesky,
          experimental_use_weight_adjustment=experimental_use_weight_adjustment,
      )
  
      samples = tf.scan(
          sampler_loop_body,
          np.arange(num_warmup_steps + num_results),
          initial_state,
      )
      return tf.nest.map_structure(
          lambda x: x[num_warmup_steps:], samples
      )
  
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 6 for calling a global function with example usage: fit_with_gibbs_sampling(model, observed_time_series, num_chains, num_results, num_warmup_steps, initial_state, seed, default_pseudo_observations, experimental_use_dynamic_cholesky, experimental_use_weight_adjustment) and returns: tf