code: def build_factored_surrogate_posterior(
      model, batch_shape=(), seed=None, name=None
  ):
      """Build a variational posterior that factors over model parameters.
  
      The surrogate posterior consists of independent Normal distributions for
      each parameter with trainable `loc` and `scale`, transformed using the
      parameter's `bijector` to the appropriate support space for that parameter.
  
      Args:
        model: An instance of `StructuralTimeSeries` representing a
            time-series model. This represents a joint distribution over
            time-series and their parameters with batch shape `[b1, ..., bN]`.
        batch_shape: Batch shape (Python `tuple`, `list`, or `int`) of initial
          states to optimize in parallel.
          Default value: `()`. (i.e., just run a single optimization).
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
        name: Python `str` name prefixed to ops created by this function.
          Default value: `None` (i.e., 'build_factored_surrogate_posterior').
      Returns:
        variational_posterior: `tfd.JointDistributionNamed` defining a trainable
            surrogate posterior over model parameters. Samples from this
            distribution are Python `dict`s with Python `str` parameter names as
            keys.
  
      ### Examples
  
      Assume we've built a structural time-series model:
  
      ```python
        day_of_week = tfp.sts.Seasonal(
            num_seasons=7,
            observed_time_series=observed_time_series,
            name='day_of_week')
        local_linear_trend = tfp.sts.LocalLinearTrend(
            observed_time_series=observed_time_series,
            name='local_linear_trend')
        model = tfp.sts.Sum(components=[day_of_week, local_linear_trend],
                            observed_time_series=observed_time_series)
      ```
  
      To fit the model to data, we define a surrogate posterior and fit it
      by optimizing a variational bound:
  
      ```python
        surrogate_posterior = tfp.sts.build_factored_surrogate_posterior(
          model=model)
        loss_curve = tfp.vi.fit_surrogate_posterior(
          target_log_prob_fn=model.joint_distribution(observed_time_series).log_prob,
          surrogate_posterior=surrogate_posterior,
          optimizer=tf_keras.optimizers.Adam(learning_rate=0.1),
          num_steps=200)
        posterior_samples = surrogate_posterior.sample(50)
  
        # In graph mode, we would need to write:
        # with tf.control_dependencies([loss_curve]):
        #   posterior_samples = surrogate_posterior.sample(50)
      ```
  
      For more control, we can also build and optimize a variational loss
      manually:
  
      ```python
        @tf.function(autograph=False)  # Ensure the loss is computed efficiently
        def loss_fn():
          return tfp.vi.monte_carlo_variational_loss(
            model.joint_distribution(observed_time_series).log_prob,
            surrogate_posterior,
            sample_size=10)
  
        optimizer = tf_keras.optimizers.Adam(learning_rate=0.1)
        for step in range(200):
          with tf.GradientTape() as tape:
            loss = loss_fn()
          grads = tape.gradient(loss, surrogate_posterior.trainable_variables)
          optimizer.apply_gradients(
            zip(grads, surrogate_posterior.trainable_variables))
          if step % 20 == 0:
            print('step {} loss {}'.format(step, loss))
  
        posterior_samples = surrogate_posterior.sample(50)
      ```
  
      """
      with tf.name_scope(name or "build_factored_surrogate_posterior"):
          prior = (
              model._joint_prior_distribution()
          )  # pylint: disable=protected-access
          batch_shape = distribution_util.expand_to_vector(
              ps.convert_to_shape_tensor(batch_shape, dtype=np.int32)
          )
          return experimental_vi.build_factored_surrogate_posterior(
              event_shape=prior.event_shape_tensor(),
              bijector=prior.experimental_default_event_space_bijector(),
              batch_shape=ps.concat(
                  [batch_shape, prior.batch_shape_tensor()], axis=0
              ),
              dtype=prior.dtype,
              seed=seed,
          )
  
file_path: probability/tensorflow_probability/python/sts/fitting.py
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental import vi
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.mcmc import dual_averaging_step_size_adaptation
  from tensorflow_probability.python.mcmc import hmc
  from tensorflow_probability.python.mcmc import sample
  from tensorflow_probability.python.mcmc import transformed_kernel
  from tensorflow_probability.python.sts.internal import util
  from tensorflow_probability.python.vi import optimization
  from tensorflow.python.util import deprecation
label_desc: Use image node_id 2 for calling a global function with example usage: build_factored_surrogate_posterior(model, batch_shape, seed, name) and returns: experimental_vi