docs: def _unflatten_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)