code: def _asvi_surrogate_for_distribution(
      dist,
      base_distribution_surrogate_fn,
      prior_substitution_rules,
      surrogate_rules,
      sample_shape=None,
  ):
      """Recursively creates ASVI surrogates, and creates new variables if needed.
  
      Args:
        dist: a `tfd.Distribution` instance.
        base_distribution_surrogate_fn: Callable to build a surrogate posterior
          for a 'base' (non-meta and non-joint) distribution, with signature
          `surrogate_posterior, variables = base_distribution_fn(
          dist, sample_shape=None, variables=None, seed=None)`.
        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 `surrogate_posterior, variables = surrogate_fn(dist,
          build_nested_surrogate_fn, sample_shape=None, variables=None, seed=None)`.
        sample_shape: Optional `Tensor` shape of samples drawn from `dist` by
          `tfd.Sample` wrappers. If not `None`, the surrogate's event will include
          independent sample dimensions, i.e., it will have event shape
          `concat([sample_shape, dist.event_shape], axis=0)`.
          Default value: `None`.
      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.
      """
      dist_name = _get_name(
          dist
      )  # Attempt to preserve the original name.
      dist = _as_substituted_distribution(
          dist, prior_substitution_rules
      )
      # Apply the first surrogate rule that matches this distribution.
      surrogate_posterior = None
      for condition, surrogate_fn in surrogate_rules:
          if _satisfies_condition(dist, condition):
              # The surrogate fn may be a generator (internal interface) or a stateless
              # trainable builder returning `(init_fn, apply_fn)` (external interface).
              maybe_gen = surrogate_fn(
                  dist,
                  build_nested_surrogate=functools.partial(
                      _asvi_surrogate_for_distribution,
                      base_distribution_surrogate_fn=base_distribution_surrogate_fn,
                      prior_substitution_rules=prior_substitution_rules,
                      surrogate_rules=surrogate_rules,
                      sample_shape=sample_shape,
                  ),
                  sample_shape=sample_shape,
              )
              if inspect.isgenerator(maybe_gen):
                  surrogate_posterior = yield from maybe_gen
              else:
                  init_fn, apply_fn = maybe_gen
                  params = yield trainable_state_util.Parameter(init_fn)
                  surrogate_posterior = apply_fn(params)
              break
      if surrogate_posterior is None:
          if (
              hasattr(dist, "distribution")
              and
              # Transformed dists not handled above are treated as base distributions.
              not isinstance(
                  dist, transformed_distribution.TransformedDistribution
              )
          ):
              raise ValueError(
                  "None of the provided substitution rules matched meta-distribution: "
                  "`{}`.".format(dist)
              )
          else:
              surrogate_posterior = (
                  yield from base_distribution_surrogate_fn(
                      dist=dist, sample_shape=sample_shape
                  )
              )
      return _set_name(surrogate_posterior, dist_name)
  
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 4 for calling a global function with example usage: _asvi_surrogate_for_distribution(dist, base_distribution_surrogate_fn, prior_substitution_rules, surrogate_rules, sample_shape) and returns: _set_name