docs: def _flatten_beam_dim(tensor):
    """Reshapes first two dimensions in to single dimension.
  
    Args:
      tensor: Tensor to reshape of shape [A, B, ...]
  
    Returns:
      Reshaped tensor of shape [A*B, ...]
    """
    shape = _shape_list(tensor)
    shape[0] *= shape[1]
    shape.pop(1)  # Remove beam dim
    return tf.reshape(tensor, shape)