code: def _windowed_adaptive_impl(
      n_draws,
      joint_dist,
      *,
      kind,
      n_chains,
      proposal_kernel_kwargs,
      num_adaptation_steps,
      current_state,
      dual_averaging_kwargs,
      trace_fn,
      return_final_kernel_results,
      discard_tuning,
      seed,
      chain_axis_names,
      **pins,
  ):
      """Runs windowed sampling using either HMC or NUTS as internal sampler."""
      if trace_fn is None:
          trace_fn = lambda *args: ()
          no_trace = True
      else:
          no_trace = False
  
      if isinstance(n_chains, int):
          n_chains = [n_chains]
  
      if (
          tf.executing_eagerly()
          or not control_flow_util.GraphOrParentsInXlaContext(
              tf1.get_default_graph()
          )
      ):
          # A Tensor num_draws argument breaks XLA, which requires static TensorArray
          # trace_fn result allocation sizes.
          num_adaptation_steps = ps.convert_to_shape_tensor(
              num_adaptation_steps
          )
  
      if "num_adaptation_steps" in dual_averaging_kwargs:
          warnings.warn(
              "Dual averaging adaptation will use the value specified in"
              " the `num_adaptation_steps` argument for its construction,"
              " hence there is no need to specify it in the"
              " `dual_averaging_kwargs` argument."
          )
  
      # TODO(b/180011931): if num_adaptation_steps is small, this throws an error.
      dual_averaging_kwargs[
          "num_adaptation_steps"
      ] = num_adaptation_steps
      dual_averaging_kwargs.setdefault(
          "reduce_fn",
          functools.partial(
              generic_math.reduce_log_harmonic_mean_exp,
              # There is only one log_accept_prob per chain, and we reduce across
              # all chains, so typically the all_gather will be gathering scalars,
              # which should be relatively efficient.
              experimental_allow_all_gather=True,
          ),
      )
      # By default, reduce over named axes for step size adaptation
      dual_averaging_kwargs.setdefault(
          "experimental_reduce_chain_axis_names", chain_axis_names
      )
      setup_seed, sample_seed = samplers.split_seed(
          samplers.sanitize_seed(seed), n=2
      )
      (
          target_log_prob_fn,
          initial_transformed_position,
          bijector,
          step_broadcast,
          batch_shape,
          shard_axis_names,
      ) = _setup_mcmc(
          joint_dist,
          n_chains=n_chains,
          init_position=current_state,
          seed=setup_seed,
          **pins,
      )
  
      if proposal_kernel_kwargs.get("step_size") is None:
          if batch_shape.shape != (
              0,
          ):  # Scalar batch has a 0-vector shape.
              raise ValueError(
                  "Batch target density must specify init_step_size. Got "
                  f"batch shape {batch_shape} from joint {joint_dist}."
              )
  
          init_step_size = _get_step_size(
              initial_transformed_position, target_log_prob_fn
          )
  
      else:
          init_step_size = step_broadcast(
              proposal_kernel_kwargs["step_size"]
          )
  
      proposal_kernel_kwargs.update(
          {
              "target_log_prob_fn": target_log_prob_fn,
              "step_size": init_step_size,
              "momentum_distribution": _init_momentum(
                  initial_transformed_position,
                  batch_shape=ps.concat(
                      [n_chains, batch_shape], axis=0
                  ),
                  shard_axis_names=shard_axis_names,
              ),
          }
      )
  
      initial_running_variance = [
          sample_stats.RunningVariance.from_stats(  # pylint: disable=g-complex-comprehension
              num_samples=tf.zeros([], part.dtype),
              mean=tf.zeros_like(part),
              variance=tf.ones_like(part),
          )
          for part in initial_transformed_position
      ]
      # TODO(phandu): Consider splitting out warmup and post warmup phases
      # to avoid executing adaptation code during the post warmup phase.
      ret = _do_sampling(
          kind=kind,
          proposal_kernel_kwargs=proposal_kernel_kwargs,
          dual_averaging_kwargs=dual_averaging_kwargs,
          num_draws=n_draws
          if discard_tuning
          else n_draws + num_adaptation_steps,
          num_burnin_steps=num_adaptation_steps
          if discard_tuning
          else 0,
          initial_position=initial_transformed_position,
          initial_running_variance=initial_running_variance,
          bijector=bijector,
          trace_fn=trace_fn,
          return_final_kernel_results=return_final_kernel_results,
          chain_axis_names=chain_axis_names,
          shard_axis_names=shard_axis_names,
          seed=sample_seed,
      )
  
      if return_final_kernel_results:
          draws, trace, fkr = ret
          return sample.CheckpointableStatesAndTrace(
              all_states=bijector.inverse(draws),
              trace=trace,
              final_kernel_results=fkr,
          )
      else:
          draws, trace = ret
          if no_trace:
              return bijector.inverse(draws)
          else:
              return sample.StatesAndTrace(
                  all_states=bijector.inverse(draws), trace=trace
              )
  
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 15 for calling a global function with example usage: _windowed_adaptive_impl(n_draws, joint_dist) and returns: sample, bijector, sample