code: def update_momentum_distribution(
      momentum_distribution, running_variance_parts
  ):
      """Updates a momentum distribution with new running variance.
  
      Args:
        momentum_distribution: Distribution arranged like a result of
          `make_momentum_distribution`.
        running_variance_parts: List of `Tensor` outputs of
          `tfp.experimental.stats.RunningVariance.variance()`.
  
      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`
      """
      model = []
      if len(running_variance_parts) != len(
          momentum_distribution.model
      ):
          raise ValueError(
              "State size mismatch: "
              f"{len(running_variance_parts)} vs {len(momentum_distribution.model)}"
          )
      for var, bb in zip(
          running_variance_parts, momentum_distribution.model
      ):
          if not isinstance(bb, batch_broadcast.BatchBroadcast):
              raise ValueError(
                  f"Part dist is not a BatchBroadcast: {bb}"
              )
          td = bb.distribution
          is_sharded, shard_axes = False, None
          if isinstance(td, sharded.Sharded):
              is_sharded, shard_axes = (
                  True,
                  td.experimental_shard_axis_names,
              )
              td = td.distribution
          if not isinstance(
              td, transformed_distribution.TransformedDistribution
          ):
              raise ValueError(
                  f"Inner dist is not a TransformedDistribution: {td}"
              )
          mvnpfl = td.distribution
          if not isinstance(
              mvnpfl,
              mvn_pfl.MultivariateNormalPrecisionFactorLinearOperator,
          ):
              raise ValueError(
                  "Inner dist is not a "
                  f"MultivariateNormalPrecisionFactorLinearOperator: {mvnpfl}"
              )
          var_flat = td.bijector.inverse(var)
          var_flat_bc = tf.broadcast_to(
              var_flat, ps.shape(mvnpfl.precision.diag_part())
          )
          mvnpfl = mvnpfl.copy(
              precision_factor=tf.linalg.LinearOperatorDiag(
                  tf.math.sqrt(var_flat_bc)
              ),
              precision=tf.linalg.LinearOperatorDiag(var_flat_bc),
          )
          td = td.copy(distribution=mvnpfl)
          if is_sharded:
              td = sharded.Sharded(td, shard_axis_name=shard_axes)
          model_dist = bb.copy(distribution=td)
          model.append(model_dist)
      return momentum_distribution.copy(model=model)
  
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 2 for calling a global function with example usage: update_momentum_distribution(momentum_distribution, running_variance_parts) and returns: momentum_distribution