PyTorch einsum with named tensors

zweisum.py

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("...")

    free_chr = ord('a')
    for name in invalid_names:
        while chr(free_chr) in names:
            free_chr += 1
        exp = exp.replace(name, chr(free_chr))
        free_chr += 1

    return exp


def zweisum(exp, *args, **kwargs):
    return torch.einsum(einsumfy_exp(exp), args)


x = torch.randn(2, 3, 5)
W = torch.randn(5, 7)

print(zweisum("batch time emb, emb hidden -> batch time hidden", x, W))