code: def _get_joint_batch_shape(batch_shape_tensor):
      if tf.nest.is_nested(batch_shape_tensor):
          return functools.reduce(
              ps.broadcast_shape, tf.nest.flatten(batch_shape_tensor)
          )
      else:
          return batch_shape_tensor
  
imports: import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import nest_util
  from tensorflow_probability.python.internal import parameter_properties
  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.mcmc.internal import util
label_desc: Use image node_id 2 for calling a global function with example usage: _get_joint_batch_shape(batch_shape_tensor) and returns: functools, batch_shape_tensor