code: def make_momentum_distribution(
      state_parts,
      batch_shape,
      running_variance_parts=None,
      shard_axis_names=None,
  ):
      """Construct a momentum distribution from the running variance.
  
      This uses a running variance to construct a momentum distribution with the
      correct batch_shape and event_shape.
  
      Args:
        state_parts: List of `Tensor`.
        batch_shape: Batch shape.
        running_variance_parts: Optional, list of `Tensor`
           outputs of `tfp.experimental.stats.RunningVariance.variance()`. Defaults
           to ones with the same shape as state_parts.
        shard_axis_names: A structure of string names indicating how members of the
          state are sharded.
  
      Returns:
        `tfd.Distribution` where `.sample` has the same structure as `state_parts`,
        and `.log_prob` of the sample will have the rank of `batch_ndims`
      """
      if running_variance_parts is None:
          running_variance_parts = tf.nest.map_structure(
              tf.ones_like, state_parts
          )
      distributions = []
      batch_ndims = ps.rank_from_shape(batch_shape)
      use_sharded_jd = True
      if shard_axis_names is None:
          use_sharded_jd = False
          shard_axis_names = [None] * len(state_parts)
      for variance_part, state_part, shard_axes in zip(
          running_variance_parts, state_parts, shard_axis_names
      ):
          event_shape = state_part.shape[batch_ndims:]
          if not tensorshape_util.is_fully_defined(event_shape):
              event_shape = ps.shape(state_part, name="state_part_shp")[
                  batch_ndims:
              ]
          variance_tiled = tf.broadcast_to(
              variance_part,
              ps.concat([batch_shape, event_shape], axis=0),
          )
          nevt = ps.cast(ps.reduce_prod(event_shape), tf.int32)
          variance_flattened = tf.reshape(
              variance_tiled, ps.concat([batch_shape, [nevt]], axis=0)
          )
  
          distribution = transformed_distribution.TransformedDistribution(
              bijector=reshape.Reshape(
                  event_shape_out=event_shape, name="reshape_mvnpfl"
              ),
              distribution=(
                  mvn_pfl.MultivariateNormalPrecisionFactorLinearOperator(
                      precision_factor=tf.linalg.LinearOperatorDiag(
                          tf.math.sqrt(variance_flattened)
                      ),
                      precision=tf.linalg.LinearOperatorDiag(
                          variance_flattened
                      ),
                      name="momentum",
                  )
              ),
          )
          if shard_axes:
              distribution = sharded.Sharded(
                  distribution, shard_axis_name=shard_axes
              )
          distributions.append(distribution)
      if use_sharded_jd:
          jd = sharded_jd_lib.JointDistributionSequential(distributions)
      else:
          jd = jds.JointDistributionSequential(distributions)
      return maybe_make_list_and_batch_broadcast(jd, batch_shape)
  
file_path: probability/tensorflow_probability/python/experimental/mcmc/preconditioning_utils.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import reshape
  from tensorflow_probability.python.distributions import batch_broadcast
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.experimental.distribute import joint_distribution
  from tensorflow_probability.python.experimental.distribute import sharded
  from tensorflow_probability.python.experimental.distributions import mvn_precision_factor_linop
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 1 for calling a global function with example usage: make_momentum_distribution(state_parts, batch_shape, running_variance_parts, shard_axis_names) and returns: maybe_make_list_and_batch_broadcast