candle-einops

This library is a fork of einops intended to bring support for einops to Candle. Thank you @VasanthakumarV for such a fantastic macro based library to build off. The original library was implemented with TCH as the backing library and was based on the einops python library.

For the most part everything from the original library has remained and only the device/dtype bindings have been altered. I do have plans to port einsum functionality to this library in the future.

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, and mean. 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")