This is a simple wrapper around np.einsum, jnp.einsum, and torch.einsum that allows the use of self-documenting variable names instead of just single letters in the equations. Inspired by the syntax in einops.
For example, instead of writing:
import torch
torch.einsum('bct,bcs->bcts', a, b)or
import numpy as np
np.einsum('bct,bcs->bcts', a, b)or
import jax.numpy as jnp
jnp.einsum('bct,bcs->bcts', a, b)With this library you can write:
from fancy_einsum import einsum
einsum('batch channels time1, batch channels time2 -> batch channels time1 time2', a, b)