import tensorflow as tf

def pad_and_gather(vecs, mask_inds, pad=None):
    """
    pad a vector with a zero row and gather with input inds
    """
    if pad is None:
        pad = tf.expand_dims(tf.zeros_like(vecs[0]), 0)
    else:
        pad = tf.expand_dims(pad, 0)
    vecs_padded = tf.concat(0, [vecs, pad])
    # flatten mask and edges
    vecs_gathered = tf.gather(vecs_padded, mask_inds)
    return vecs_gathered

def padded_segment_reduce(vecs, segment_inds, num_segments, reduction_mode):
    """
    Reduce the vecs with segment_inds and reduction_mode
    Input:
        vecs: A Tensor of shape (batch_size, vec_dim)
        segment_inds: A Tensor containing the segment index of each
        vec row, should agree with vecs in shape[0]
