inox.nn.einops

Einstein-like tensor operations

References

Einops: Clear and Reliable Tensor Manipulations with Einstein-like Notation (Rogozhnikov et al., 2022)
https://openreview.net/forum?id=oapKSVM2bcj

Classes

Rearrange

Creates an axis rearrangement layer.

Reduce

Creates an axis reduction layer.

Repeat

Creates an axis repetition layer.

Descriptions

class inox.nn.einops.Rearrange(pattern, **lengths)

Creates an axis rearrangement layer.

This module is a thin wrapper around einops.rearrange.

Parameters:

  • pattern (str) – The axis rearrangement pattern. For example, the pattern 'A B C -> C (A B)' moves and flattens the two first axes.

  • lengths (int) – The lengths of the axes.

class inox.nn.einops.Reduce(pattern, reduction='sum', **lengths)

Creates an axis reduction layer.

This module is a thin wrapper around einops.reduce.

Parameters:

  • pattern (str) – The axis rearrangement pattern. For example, the pattern 'A B C -> A C' reduces the second axis.

  • reduction (str) – The type of reduction ('sum', 'mean', 'max', …).

  • lengths (int) – The lengths of the axes.

class inox.nn.einops.Repeat(pattern, **lengths)

Creates an axis repetition layer.

This module is a thin wrapper around einops.repeat.

Parameters:

  • pattern (str) – The axis rearrangement pattern. For example, the pattern 'A B -> A C B' inserts a new axis.

  • lengths (int) – The lengths of the axes.