code: def maybe_make_list_and_batch_broadcast(
      momentum_distribution, batch_shape
  ):
      """Makes the distribution list-like and batched, if possible."""
      if not mcmc_util.is_list_like(momentum_distribution.dtype):
          momentum_distribution = jds.JointDistributionSequential(
              [momentum_distribution], name="joint_momentum"
          )
      if (
          isinstance(
              momentum_distribution, jds.JointDistributionSequential
          )
          and not isinstance(
              momentum_distribution, jdn.JointDistributionNamed
          )
          and
          # Skip this step if we already batch broadcast.
          not all(
              isinstance(md, batch_broadcast.BatchBroadcast)
              for md in momentum_distribution.model
          )
          and not any(
              callable(dist_fn)
              for dist_fn in momentum_distribution.model
          )
      ):
          momentum_distribution = momentum_distribution.copy(
              model=[
                  batch_broadcast.BatchBroadcast(
                      md, with_shape=batch_shape
                  )
                  for md in momentum_distribution.model
              ]
          )
      return momentum_distribution
  
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 3 for calling a global function with example usage: maybe_make_list_and_batch_broadcast(momentum_distribution, batch_shape) and returns: momentum_distribution