einops.einsum

Hey! Loving einops, so much that now I feel a bit sad about standard einsum not being able to use descriptive names for dimensions. It would be amazing if einops implemented einsum with the same conveniences.

[cgarciae]

Hey @arogozhnikov, I’d already seen the WeightedEinsum layer, its very nice! I still think it would useful to have the functional version for the operation.

[MilesCranmer]

+1 to this! My code right now is a mix of einops and torch.einsum, but I think it would be very nice/consistent if it was all the same syntax through einops. As the name einops suggests, I half expected einsum to be included when I tried it the first time.

I think a wrapper could start off by just being something that takes, e.g.,

einops.einsum(x, y, 'time i j, time i k -> time i k')

and simply parses it into usual einsum notation (for each corresponding backend):

torch.einsum("aij,ajk->aik", x, y)

Maybe eventually there could be cool features to consider related to the other syntax innovations in this awesome package.

Cheers! Miles