TensorFlow utility method to obtain the correct last outputs from a tf.nn.dynamic_rnn as determined by sequence length

gather_by_lengths.py

def gather_by_lengths(outputs, seq_lengths):
    """
    :param outputs: [batch_size x max_seq_length x hidden_size] tensor of dynamic_rnn outputs
    :param seq_lengths: [batch_size] tensor of sequence lengths
    :return: [batch_size x hidden_size] tensor of last outputs
    """
    batch_size, max_seq_length, hidden_size = tf.unpack(tf.shape(outputs))
    index = tf.range(0, batch_size) * max_seq_length + (seq_lengths - 1)
    return tf.gather(tf.reshape(outputs, [-1, hidden_size]), index)