code: def _setup_mcmc(
      model, n_chains, *, init_position=None, seed=None, **pins
  ):
      """Construct bijector and transforms needed for windowed MCMC.
  
      This pins the initial model, constructs a bijector that unconstrains and
      flattens each dimension and adds a leading batch shape of `n_chains`,
      initializes a point in the unconstrained space, and constructs a transformed
      log probability using the bijector.
  
      Note that we must manually construct this target log probability instead of
      using a transformed transition kernel because the TTK assumes the shape
      in is the same as the shape out.
  
      Args:
        model: `tfd.JointDistribution`
          The model to sample from.
        n_chains: list of ints
          Number of chains (independent examples) to run.
        init_position: Optional
          Structure of tensors at which to initialize sampling. Should have the
          same shape and structure as
          `model.experimental_pin(**pins).sample_unpinned(n_chains)`.
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
        **pins:
          Values passed to `model.experimental_pin`.
  
  
      Returns:
        target_log_prob_fn: Callable on the transformed space.
        initial_transformed_position: `tf.Tensor`, sampled from a uniform (-2, 2).
        bijector: `tfb.Bijector` instance, which unconstrains and flattens.
        step_broadcast_fn: Callable to broadcast step size over latent structure.
        batch_shape: Batch shape of the model.
        shard_axis_names: Shard axis names for the model
      """
      pinned_model = model.experimental_pin(**pins) if pins else model
      bijector, step_bijector = _get_flat_unconstraining_bijector(
          pinned_model
      )
  
      if init_position is None:
          raw_init_dist = initialization.init_near_unconstrained_zero(
              pinned_model
          )
          init_position = initialization.retry_init(
              raw_init_dist.sample,
              target_fn=pinned_model.unnormalized_log_prob,
              sample_shape=n_chains,
              seed=seed,
          )
  
      initial_transformed_position = tf.nest.map_structure(
          tf.identity, bijector.forward(init_position)
      )
  
      batch_shape = pinned_model.batch_shape
      if tf.nest.is_nested(batch_shape):
          batch_shape = functools.reduce(
              tf.broadcast_static_shape, tf.nest.flatten(batch_shape)
          )
  
      if not tensorshape_util.is_fully_defined(batch_shape):
          batch_shape = pinned_model.batch_shape_tensor()
          if tf.nest.is_nested(batch_shape):
              batch_shape = functools.reduce(
                  tf.broadcast_dynamic_shape,
                  tf.nest.flatten(batch_shape),
              )
  
      # This tf.function is not redundant with the ones on _fast_window
      # and _slow_window because the various kernels (like HMC) may invoke
      # `target_log_prob_fn` multiple times within one window.
      @tf.function(autograph=False)
      def target_log_prob_fn(*args):
          lp = pinned_model.unnormalized_log_prob(
              bijector.inverse(args)
          )
          ldj = bijector.inverse_log_det_jacobian(
              args,
              event_ndims=[1 for _ in initial_transformed_position],
          )
          return lp + ldj
  
      def step_broadcast(step_size):
          # Only apply the bijector to nested step sizes or non-scalar batches.
          if tf.nest.is_nested(step_size):
              return step_bijector(
                  nest_util.broadcast_structure(
                      pinned_model.event_shape_tensor(), step_size
                  )
              )
          else:
              return step_size
  
      shard_axis_names = pinned_model.experimental_shard_axis_names
      if any(tf.nest.flatten(shard_axis_names)):
          shard_axis_names = nest.flatten_up_to(
              initial_transformed_position,
              list(pinned_model._model_flatten(shard_axis_names)),
          )  # pylint: disable=protected-access
  
      else:
          # No active shard axis names
          shard_axis_names = None
  
      return (
          target_log_prob_fn,
          initial_transformed_position,
          bijector,
          step_broadcast,
          ps.convert_to_shape_tensor(batch_shape, name="batch_shape"),
          shard_axis_names,
      )
  
imports: import collections
  import functools
  import warnings
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.bijectors import joint_map
  from tensorflow_probability.python.bijectors import reshape
  from tensorflow_probability.python.bijectors import restructure
  from tensorflow_probability.python.experimental.mcmc import diagonal_mass_matrix_adaptation
  from tensorflow_probability.python.experimental.mcmc import initialization
  from tensorflow_probability.python.experimental.mcmc import preconditioned_hmc
  from tensorflow_probability.python.experimental.mcmc import preconditioned_nuts
  from tensorflow_probability.python.experimental.mcmc import preconditioning_utils
  from tensorflow_probability.python.experimental.mcmc import sharded
  from tensorflow_probability.python.experimental.stats import sample_stats
  from tensorflow_probability.python.internal import nest_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.internal import unnest
  from tensorflow_probability.python.math import generic
  from tensorflow_probability.python.mcmc import dual_averaging_step_size_adaptation
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import sample
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow.python.ops import control_flow_util
  from tensorflow.python.util import nest
label_desc: Use image node_id 4 for calling a global function with example usage: _setup_mcmc(model, n_chains) and returns: target_log_prob_fn,  initial_transformed_position,  bijector,  step_broadcast,  ps,  shard_axis_names, unknown, step_bijector, step_size