code: def windowed_adaptive_nuts(
      n_draws,
      joint_dist,
      *,
      n_chains=64,
      num_adaptation_steps=500,
      current_state=None,
      init_step_size=None,
      dual_averaging_kwargs=None,
      max_tree_depth=10,
      max_energy_diff=500.0,
      unrolled_leapfrog_steps=1,
      parallel_iterations=10,
      trace_fn=default_nuts_trace_fn,
      return_final_kernel_results=False,
      discard_tuning=True,
      chain_axis_names=None,
      seed=None,
      **pins,
  ):
      """Adapt and sample from a joint distribution using NUTS, conditioned on pins.
  
      Step size is tuned using a dual-averaging adaptation, and the kernel is
      conditioned using a diagonal mass matrix, which is estimated using expanding
      windows.
  
      Args:
        n_draws: int
          Number of draws after adaptation.
        joint_dist: `tfd.JointDistribution`
          A joint distribution to sample from.
        n_chains: int or list of ints
          Number of independent chains to run MCMC with.
        num_adaptation_steps: int
          Number of draws used to adapt step size and mass matrix.
        current_state: Optional
          Structure of tensors at which to initialize sampling. Should have the
          same shape and structure as
          `model.experimental_pin(**pins).sample(n_chains)`.
        init_step_size: Optional
          Where to initialize the step size for the leapfrog integrator. The
          structure should broadcast with `current_state`. For example, if the
          initial state is
          ```
          {'a': tf.zeros(n_chains),
           'b': tf.zeros([n_chains, n_features])}
           ```
          then any of `1.`, `{'a': 1., 'b': 1.}`, or
          `{'a': tf.ones(n_chains), 'b': tf.ones([n_chains, n_features])}` will
          work. Defaults to the dimension of the log density to the 0.25 power.
        dual_averaging_kwargs: Optional dict
          Keyword arguments to pass to `tfp.mcmc.DualAveragingStepSizeAdaptation`.
          By default, a `target_accept_prob` of 0.85 is set, acceptance
          probabilities across chains are reduced using a harmonic mean, and the
          class defaults are used otherwise.
        max_tree_depth: Maximum depth of the tree implicitly built by NUTS. The
          maximum number of leapfrog steps is bounded by `2**max_tree_depth` i.e.
          the number of nodes in a binary tree `max_tree_depth` nodes deep. The
          default setting of 10 takes up to 1024 leapfrog steps.
        max_energy_diff: Scalar threshold of energy differences at each leapfrog,
          divergence samples are defined as leapfrog steps that exceed this
          threshold. Default to 1000.
        unrolled_leapfrog_steps: The number of leapfrogs to unroll per tree
          expansion step. Applies a direct linear multipler to the maximum
          trajectory length implied by max_tree_depth. Defaults to 1.
        parallel_iterations: The number of iterations allowed to run in parallel.
          It must be a positive integer. See `tf.while_loop` for more details.
        trace_fn: Optional callable
          The trace function should accept the arguments
          `(state, bijector, is_adapting, phmc_kernel_results)`,  where the `state`
          is an unconstrained, flattened float tensor, `bijector` is the
          `tfb.Bijector` that is used for unconstraining and flattening,
          `is_adapting` is a boolean to mark whether the draw is from an adaptation
          step, and `phmc_kernel_results` is the
          `UncalibratedPreconditionedHamiltonianMonteCarloKernelResults` from the
          `PreconditionedHamiltonianMonteCarlo` kernel. Note that
          `bijector.inverse(state)` will provide access to the current draw in the
          untransformed space, using the structure of the provided `joint_dist`.
        return_final_kernel_results: If `True`, then the final kernel results are
          returned alongside the chain state and the trace specified by the
          `trace_fn`.
        discard_tuning: bool
          Whether to return tuning traces and draws.
        chain_axis_names: A `str` or list of `str`s indicating the named axes
          by which multiple chains are sharded. See `tfp.experimental.mcmc.Sharded`
          for more context.
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
        **pins:
          These are used to condition the provided joint distribution, and are
          passed directly to `joint_dist.experimental_pin(**pins)`.
      Returns:
        A single structure of draws is returned in case the trace_fn is `None`, and
        `return_final_kernel_results` is `False`. If there is a trace function,
        the return value is a tuple, with the trace second. If the
        `return_final_kernel_results` is `True`, the return value is a tuple of
        length 3, with final kernel results returned last. If `discard_tuning` is
        `True`, the tensors in `draws` and `trace` will have length `n_draws`,
        otherwise, they will have length `n_draws + num_adaptation_steps`.
      """
      if dual_averaging_kwargs is None:
          dual_averaging_kwargs = {}
      dual_averaging_kwargs = dict(dual_averaging_kwargs)
      dual_averaging_kwargs.setdefault("target_accept_prob", 0.85)
      proposal_kernel_kwargs = {
          "step_size": init_step_size,
          "max_tree_depth": max_tree_depth,
          "max_energy_diff": max_energy_diff,
          "unrolled_leapfrog_steps": unrolled_leapfrog_steps,
          "parallel_iterations": parallel_iterations,
      }
      return _windowed_adaptive_impl(
          n_draws=n_draws,
          joint_dist=joint_dist,
          kind="nuts",
          n_chains=n_chains,
          proposal_kernel_kwargs=proposal_kernel_kwargs,
          current_state=current_state,
          num_adaptation_steps=num_adaptation_steps,
          dual_averaging_kwargs=dual_averaging_kwargs,
          trace_fn=trace_fn,
          return_final_kernel_results=return_final_kernel_results,
          discard_tuning=discard_tuning,
          chain_axis_names=chain_axis_names,
          seed=seed,
          **pins,
      )
  
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 13 for calling a global function with example usage: windowed_adaptive_nuts(n_draws, joint_dist) and returns: _windowed_adaptive_impl