code: def _asvi_surrogate_for_joint_distribution(
      dist, build_nested_surrogate
  ):
      """Builds a structured joint surrogate posterior for a joint model."""
  
      # Probabilistic program for ASVI surrogate posterior.
      prior_coroutine = (
          dist._model_coroutine
      )  # pylint: disable=protected-access
  
      def posterior_generator():
          # Run the prior model coroutine and create a surrogate for each
          # distribution in turn. This coroutine yields *both* the surrogate
          # parameters and the surrogate distributions themselves. It can therefore
          # run in a `trainable_state_util` context to define parameters (using
          # the `_get_coroutine_parameters` wrapper to hide the yielded distributions)
          # *and* as a JDCoroutine model (using the `_set_coroutine_parameters`
          # wrapper to hide the yielded parameters).
          prior_gen = prior_coroutine()
          dist = next(prior_gen)
          i = 0
          try:
              while True:
                  surrogate_posterior = (
                      yield from build_nested_surrogate(
                          dist.distribution
                          if isinstance(dist, Root)
                          else dist
                      )
                  )
                  if isinstance(dist, Root):
                      surrogate_posterior = Root(surrogate_posterior)
                  value_out = yield surrogate_posterior
                  dist = prior_gen.send(value_out)
                  i += 1
          except StopIteration:
              pass
  
      parameters = yield from _get_coroutine_parameters(
          posterior_generator, seed=samplers.zeros_seed()
      )
      surrogate_posterior = joint_distribution_auto_batched.JointDistributionCoroutineAutoBatched(
          _set_coroutine_parameters(posterior_generator, parameters)
      )
  
      # Ensure that the surrogate posterior structure matches that of the prior.
      try:
          tf.nest.assert_same_structure(
              dist.dtype, surrogate_posterior.dtype
          )
      except TypeError:
          tokenize = lambda jd: jd._model_unflatten(  # pylint: disable=protected-access, g-long-lambda
              range(
                  len(jd._model_flatten(jd.dtype))
              )  # pylint: disable=protected-access
          )
          surrogate_posterior = restructure.Restructure(
              output_structure=tokenize(dist),
              input_structure=tokenize(surrogate_posterior),
          )(surrogate_posterior)
      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 8 for calling a global function with example usage: _asvi_surrogate_for_joint_distribution(dist, build_nested_surrogate) and returns: surrogate_posterior