code: def _asvi_surrogate_for_markov_chain(dist, build_nested_surrogate):
      """Builds a structured surrogate posterior for a Markov chain."""
      surrogate_prior = yield from build_nested_surrogate(
          dist.initial_state_prior
      )
  
      (
          transition_all_steps_init_fn,
          _,
      ) = trainable_state_util.as_stateless_builder(
          lambda: build_nested_surrogate(  # pylint: disable=g-long-lambda
              dist.transition_fn(
                  tf.range(dist.num_steps - 1),
                  dist.initial_state_prior.sample(
                      dist.num_steps - 1, seed=samplers.zeros_seed()
                  ),
              )
          )
      )()
      transition_params = yield trainable_state_util.Parameter(
          transition_all_steps_init_fn,
          name="markov_chain_transition_params",
      )
  
      build_transition_one_step = trainable_state_util.as_stateless_builder(
          lambda step, state: build_nested_surrogate(  # pylint: disable=g-long-lambda
              dist.transition_fn(step, state)
          )
      )
  
      def surrogate_transition_fn(step, state):
          _, one_step_apply_fn = build_transition_one_step(step, state)
          return one_step_apply_fn(
              tf.nest.map_structure(
                  # Gather parameters for this specific step of the chain.
                  lambda v: tf.gather(v, step, axis=0),
                  transition_params,
              )
          )
  
      return markov_chain.MarkovChain(
          initial_state_prior=surrogate_prior,
          transition_fn=surrogate_transition_fn,
          num_steps=dist.num_steps,
          validate_args=dist.validate_args,
      )
  
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 9 for calling a global function with example usage: _asvi_surrogate_for_markov_chain(dist, build_nested_surrogate) and returns: markov_chain, one_step_apply_fn