Einops support

[turian]

Is your feature request related to a problem? Please describe.

When using einops operations, it would be great to decorate the nodes with the einops names.
It makes it much easier to keep track of which dimension means what.

Describe the solution you'd like

Rearrange('(B N) K T -> (B N) T K'

would show

((B N)=384, K=259, T=41)
----------------------------
(B=3, N=259, T=41)

in the right side of the displayed cell:

Describe alternatives you've considered

Spending more time looking at torchview :)

[mert-kurttutan]

One way I can think of to integrate this feature is to use self.pattern for einops modules since its object structure is

class ReduceMixin:
    """
    Reduce layer behaves identically to einops.reduce operation.
    :param pattern: str, rearrangement pattern
    :param reduction: one of available reductions ('min', 'max', 'sum', 'mean', 'prod'), case-sensitive
    :param axes_lengths: any additional specification of dimensions
    See einops.reduce for source_examples.
    """

    def __init__(self, pattern: str, reduction: str, **axes_lengths: Any):
        super().__init__()
        self.pattern = pattern
        self.reduction = reduction
        self.axes_lengths = axes_lengths
        self._recipe = self.recipe()  # checking parameters

One should think about the dependency if this feature were to be integrated. I think one way to go is to create an optional import and then check if module is Rearrange from einops. If the module is imported successfully, then use isinstance(module, einops.Rearrange) to include self.pattern in the name string. Otherwise, don't do this.

In the meantime, let me think about the cons and pros of adding this feature, if there is any