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
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")