Oven logo

Oven

Published

Drop-in replacement for torch/numpy einsum, with descriptive variable names in equations

pip install fancy-einsum

Package Downloads

Weekly DownloadsMonthly Downloads

Requires Python

>=3.6

Dependencies

    Fancy Einsum

    This is a simple wrapper around np.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)
    

    With this library you can write:

    from fancy_einsum import einsum
    einsum('batch channels time1, batch channels time2 -> batch channels time1 time2', a, b)