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.