code: def hmc_like_momentum_distribution_setter_fn(
      kernel_results, new_distribution
  ):
      """Setter for `momentum_distribution` so it can be adapted."""
      # Note that unnest.replace_innermost has a special path for going into
      # `accepted_results` preferentially, so this will set
      # `accepted_results.momentum_distribution`.
      return unnest.replace_innermost(
          kernel_results, momentum_distribution=new_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 1 for calling a global function with example usage: hmc_like_momentum_distribution_setter_fn(kernel_results, new_distribution) and returns: unnest