einops

This library is heavily inspired by python's einops.

Currently tch is the only available backend.

Difference from the python version,

• All code generated at compile time, avoiding the need for caching
• One common api for rearrange, reduce and repeat operations
• Shape and reduction operations can be directly specified in the expression

Getting Started

Transpose

Permute/Transpose dimensions, left side of -> is the original state, right of -> describes the end state

// (28, 28, 3) becomes (3, 28, 28)
let output = einops!("h w c -> c h w", &input);

Composition

Combine dimensions by putting them inside a parenthesis on the right of ->

// (10, 28, 28, 3) becomes (280, 28, 3)
let output = einops!("b h w c -> (b h) w c", &input);

Transpose + Composition

Transpose a tensor, followed by a composing two dimensions into one, in one single expression

// (10, 28, 28, 3) becomes (28, 280, 3)
let output = einops!("b h w c -> h (b w) c", &input);

Decomposition

Split a dimension into two, by specifying the details inside parenthesis on the left, specify the shape of the new dimensions like so b1:2, b1 is a new dimension with shape 2

// (10, 28, 28, 3) becomes (2, 5, 28, 28, 3)
let output = einops!("(b1:2 b2) h w c -> b1 b2 h w c", &input);

New axis can also be specified from variables or fields (struct and enum) using curly braces

let b1 = 2;
let output = einops!("({b1} b2) h w c -> {b1} b2 h w c", &input);

Decomposition + Transpose + Composition

We can perform all operations discussed so far in a single expression

// (10, 28, 28, 3) becomes (56, 140 3)
let output = einops!("b h (w w2:2) c -> (h w2) (b w) c", &input);

Reduce

We can reduce axes using operations like, sum, min, max, mean and prod, if the same operations has to be performed on multiple continuous axes we can do sum(a b c)

// (10, 28, 28, 3) becomes (28, 28, 3)
let output = einops!("mean(b) h w c -> h w c", &input);

Decomposition + Reduce + Transpose + Composition

Single expression for combining all functionalities discussed

// (10, 28, 28, 3) becomes (14, 140, 3)
let output = einops!("b (h max(h2:2)) (w max(w2:2)) c -> h (b w) c", &input);

Repeat

We can repeat axes by specify it on the right side of ->, it can named, or it can simply be a number

// (28, 28, 3) becomes (28, 5, 28, 3)
let output = einops!("h w c -> h repeat:5 w c", &input);

Repeating axis's shape can be from a variables or a field (struct, enum)

let repeat = 5;
let output = einops!("h w c -> h {repeat} w c", &input);

Squeeze

Squeeze axes of shape 1

// (1, 28, 28, 3) becomes (28, 28, 3)
let output = einops!("1 h w c -> h w c")