code: def _init_momentum(
      initial_transformed_position, *, batch_shape, shard_axis_names
  ):
      """Initialize momentum so trace_fn can be concatenated."""
      variance_parts = [
          ps.ones_like(p) for p in initial_transformed_position
      ]
      return preconditioning_utils.make_momentum_distribution(
          state_parts=initial_transformed_position,
          batch_shape=batch_shape,
          running_variance_parts=variance_parts,
          shard_axis_names=shard_axis_names,
      )
  
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 12 for calling a global function with example usage: _init_momentum(initial_transformed_position) and returns: preconditioning_utils