code: def _make_flatten_unflatten_fns(batch_shape):
      """Builds functions for flattening and unflattening batch dimensions."""
      batch_shape = tuple(batch_shape)
      batch_rank = len(batch_shape)
      ndims = ps.cast(ps.reduce_prod(batch_shape), tf.int32)
  
      def flatten_fn(x):
          x_shape = tuple(x.shape)
          if x_shape[:batch_rank] != batch_shape:
              raise ValueError(
                  "Expected batch-shape=%s; received array of shape=%s"
                  % (batch_shape, x_shape)
              )
          flat_shape = (ndims,) + x_shape[batch_rank:]
          return tf.reshape(x, flat_shape)
  
      def unflatten_fn(x):
          x_shape = tuple(x.shape)
          if x_shape[0] != ndims:
              raise ValueError(
                  "Expected batch-size=%d; received shape=%s"
                  % (ndims, x_shape)
              )
          return tf.reshape(x, batch_shape + x_shape[1:])
  
      return flatten_fn, unflatten_fn
  
file_path: probability/tensorflow_probability/python/stats/calibration.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import generic
  from tensorflow_probability.python.stats import quantiles
label_desc: Use image node_id 5 for calling a global function with example usage: _make_flatten_unflatten_fns(batch_shape) and returns: flatten_fn,  unflatten_fn, tf, tf