5 unstable releases
0.3.0-alpha.2 | Jun 23, 2022 |
---|---|
0.3.0-alpha.1 | Dec 31, 2021 |
0.2.1 | Sep 4, 2021 |
0.2.0 | Sep 4, 2021 |
0.1.0 | Apr 7, 2021 |
#340 in Machine learning
33 downloads per month
16KB
248 lines
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")
Dependencies
~1.2–4MB
~75K SLoC