rockt

import torch
import re

def einsumfy_exp(exp):
    names = set(re.split("[, \(\)]|->", exp))
    names.remove("")

    invalid_names = set(filter(lambda x: len(x) > 1, names))
    if "..." in invalid_names:
        invalid_names.remove("...")

rockt

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)