code: def default_nuts_trace_fn(state, bijector, is_adapting, pkr):
      """Trace function for `windowed_adaptive_nuts` providing standard diagnostics.
  
      Specifically, these match up with a number of diagnostics used by ArviZ [1],
      to make diagnostics and analysis easier. The names used follow those used in
      TensorFlow Probability, and will need to be mapped to those used in the ArviZ
      schema.
  
      References:
        [1]: Kumar, R., Carroll, C., Hartikainen, A., & Martin, O. (2019). ArviZ a
        unified library for exploratory analysis of Bayesian models in Python.
        Journal of Open Source Software, 4(33), 1143.
  
      Args:
       state: tf.Tensor
         Current sampler state, flattened and unconstrained.
       bijector: tfb.Bijector
         This can be used to untransform the shape to something with the same shape
         as will be returned.
       is_adapting: bool
         Whether this is an adapting step, or may be treated as a valid MCMC draw.
       pkr: UncalibratedPreconditionedHamiltonianMonteCarloKernelResults
         Kernel results from this iteration.
  
      Returns:
        dict with sampler statistics.
      """
      del state, bijector  # Unused
  
      energy_diff = unnest.get_innermost(pkr, "log_accept_ratio")
      return {
          "step_size": unnest.get_innermost(pkr, "step_size"),
          "tune": is_adapting,
          "target_log_prob": unnest.get_innermost(
              pkr, "target_log_prob"
          ),
          "diverging": unnest.get_innermost(pkr, "has_divergence"),
          "accept_ratio": tf.minimum(
              tf.ones_like(energy_diff), tf.exp(energy_diff)
          ),
          "variance_scaling": unnest.get_innermost(
              pkr, "momentum_distribution"
          ).variance(),
          "n_steps": unnest.get_innermost(pkr, "leapfrogs_taken"),
          "is_accepted": unnest.get_innermost(pkr, "is_accepted"),
      }
  
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 1 for calling a global function with example usage: default_nuts_trace_fn(state, bijector, is_adapting, pkr) and returns: dict