code: def _build_asvi_surrogate_posterior(
      prior,
      mean_field=False,
      initial_prior_weight=0.5,
      prior_substitution_rules=(ASVI_DEFAULT_PRIOR_SUBSTITUTION_RULES),
      surrogate_rules=(ASVI_DEFAULT_SURROGATE_RULES),
      name=None,
  ):
      """Builds a structured surrogate posterior inspired by conjugate updating.
  
      ASVI, or Automatic Structured Variational Inference, was proposed by
      Ambrogioni et al. (2020) [1] as a method of automatically constructing a
      surrogate posterior with the same structure as the prior. It does this by
      reparameterizing the variational family of the surrogate posterior by
      structuring each parameter according to the equation
      ```none
      prior_weight * prior_parameter + (1 - prior_weight) * mean_field_parameter
      ```
      In this equation, `prior_parameter` is a vector of prior parameters and
      `mean_field_parameter` is a vector of trainable parameters with the same
      domain as `prior_parameter`. `prior_weight` is a vector of learnable
      parameters where `0. <= prior_weight <= 1.`. When `prior_weight =
      0`, the surrogate posterior will be a mean-field surrogate, and when
      `prior_weight = 1.`, the surrogate posterior will be the prior. This convex
      combination equation, inspired by conjugacy in exponential families, thus
      allows the surrogate posterior to balance between the structure of the prior
      and the structure of a mean-field approximation.
  
      Args:
        prior: tfd.JointDistribution instance of the prior.
        mean_field: Optional Python boolean. If `True`, creates a degenerate
          surrogate distribution in which all variables are independent,
          ignoring the prior dependence structure. Default value: `False`.
        initial_prior_weight: Optional float value (either static or tensor value)
          on the interval [0, 1]. A larger value creates an initial surrogate
          distribution with more dependence on the prior structure. Default value:
          `0.5`.
        prior_substitution_rules: Iterable of substitution rules applied to the
          prior before constructing a surrogate. Each rule is a `(condition,
          substitution_fn)` tuple; these are checked in order and *all* applicable
          substitutions are made. The `condition` may be either a class or a
          callable returning a boolean (for example, `tfd.Normal` or, equivalently,
          `lambda dist: isinstance(dist, tfd.Normal)`). The `substitution_fn` should
          have signature `new_dist = substitution_fn(dist)`.
        surrogate_rules: Iterable of special-purpose rules to create surrogates
          for specific distribution types. Each rule is a `(condition,
          surrogate_fn)` tuple; these are checked in order and the first applicable
          `surrogate_fn` is used. The `condition` may be either a class or a
          callable returning a boolean (for example, `tfd.Normal` or, equivalently,
          `lambda dist: isinstance(dist, tfd.Normal)`). The `surrogate_fn` should
          have signature `init_fn, apply_fn = surrogate_fn(dist,
          build_nested_surrogate_fn, sample_shape=None)`.
        name: Optional string. Default value: `build_asvi_surrogate_posterior`.
      Yields:
        *parameters: sequence of `trainable_state_util.Parameter` namedtuples.
          These are intended to be consumed by
          `trainable_state_util.as_stateful_builder` and
          `trainable_state_util.as_stateless_builder` to define stateful and
          stateless variants respectively.
  
      ### Examples
  
      Consider a Brownian motion model expressed as a JointDistribution:
  
      ```python
      prior_loc = 0.
      innovation_noise = .1
  
      def model_fn():
        new = yield tfd.Normal(loc=prior_loc, scale=innovation_noise)
        for i in range(4):
          new = yield tfd.Normal(loc=new, scale=innovation_noise)
  
      prior = tfd.JointDistributionCoroutineAutoBatched(model_fn)
      ```
  
      Let's use variational inference to approximate the posterior. We'll build a
      surrogate posterior distribution by feeding in the prior distribution.
  
      ```python
      surrogate_posterior =
        tfp.experimental.vi.build_asvi_surrogate_posterior(prior)
      ```
  
      This creates a trainable joint distribution, defined by variables in
      `surrogate_posterior.trainable_variables`. We use `fit_surrogate_posterior`
      to fit this distribution by minimizing a divergence to the true posterior.
  
      ```python
      losses = tfp.vi.fit_surrogate_posterior(
        target_log_prob_fn,
        surrogate_posterior=surrogate_posterior,
        num_steps=100,
        optimizer=tf_keras.optimizers.Adam(0.1),
        sample_size=10)
  
      # After optimization, samples from the surrogate will approximate
      # samples from the true posterior.
      samples = surrogate_posterior.sample(100)
      posterior_mean = [tf.reduce_mean(x) for x in samples]
      posterior_std = [tf.math.reduce_std(x) for x in samples]
      ```
  
      #### References
  
      [1]: Luca Ambrogioni, Kate Lin, Emily Fertig, Sharad Vikram, Max Hinne,
            Dave Moore, Marcel van Gerven. Automatic structured variational
            inference. _arXiv preprint arXiv:2002.00643_, 2020
            https://arxiv.org/abs/2002.00643
  
      """
      with tf.name_scope(name or "build_asvi_surrogate_posterior"):
          surrogate_posterior = (
              yield from _asvi_surrogate_for_distribution(
                  dist=prior,
                  base_distribution_surrogate_fn=functools.partial(
                      _asvi_convex_update_for_base_distribution,
                      mean_field=mean_field,
                      initial_prior_weight=initial_prior_weight,
                  ),
                  prior_substitution_rules=prior_substitution_rules,
                  surrogate_rules=surrogate_rules,
              )
          )
          return surrogate_posterior
  
imports: import collections
  import copy
  import functools
  import inspect
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import restructure
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.bijectors import sigmoid
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import chi2
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import gamma
  from tensorflow_probability.python.distributions import half_normal
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import markov_chain
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.distributions import truncated_normal
  from tensorflow_probability.python.distributions import uniform
  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 samplers
  from tensorflow_probability.python.internal import trainable_state_util
label_desc: Use image node_id 11 for calling a global function with example usage: _build_asvi_surrogate_posterior(prior, mean_field, initial_prior_weight, prior_substitution_rules, surrogate_rules, name) and returns: surrogate_posterior