Skip to content

VasanthakumarV/einops

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

einops crates docs

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