code: def hmc_like_momentum_distribution_getter_fn(kernel_results):
      """Getter for `momentum_distribution` so it can be updated."""
      return unnest.get_innermost(
          kernel_results, "momentum_distribution"
      )
  
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.mcmc import preconditioning_utils
  from tensorflow_probability.python.experimental.stats import sample_stats
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import unnest
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 2 for calling a global function with example usage: hmc_like_momentum_distribution_getter_fn(kernel_results) and returns: unnest