code: def _unmerge_beam_dim(tensor, batch_size, beam_size):
      """Reshapes first dimension back to [batch_size, beam_size].
  
      Args:
        tensor: Tensor to reshape of shape [batch_size*beam_size, ...]
        batch_size: Tensor, original batch size.
        beam_size: int, original beam size.
  
      Returns:
        Reshaped tensor of shape [batch_size, beam_size, ...]
      """
      shape = shape_list(tensor)
      new_shape = [batch_size] + [beam_size] + shape[1:]
      return tf.reshape(tensor, new_shape)
  
file_path: finetune/finetune/util/beam_search.py
imports: import math
  import numpy
  from finetune.util.shapes import shape_list
  import tensorflow
  from tensorflow.python.ops import inplace_ops
  from tensorflow.python.util import nest
label_desc: Use image node_id 5 for calling a global function with example usage: _unmerge_beam_dim(tensor, batch_size, beam_size) and returns: tf