Crate einops[−][src]
Einops
This is a rust port of the incredible einops library. Almost all the operations specified in its tutorial should be available, if you find any inconsistencies please raise a github issue.
Unlike its python counterpart, caching the parsed expression has not been implemented yet. So
when applying the same pattern multiple times, prefer Rearrange::new(...)
or Rearrange::with_lengths(...)
api, over the methods available through RearrangeFn
like traits
Flexible and powerful tensor operations for readable and reliable code. Currently only supports tch.
Getting started
Add the following to your Cargo.toml
file,
[dependencies] einops = { version: "0.1", features: ["tch-bindings"] }
Examples
Einops provies three operations, they cover stacking, reshape, transposition, squeeze/unsqueeze, repeat, tile, concatenate and numerous reductions
// Tch specific imports use tch::{Tensor, Kind, Device}; // Structs that provide constructor like api use einops::{Rearrange, Repeat, Reduce, Operation, Backend}; // Traits required to call functions directly on the tensors use einops::{ReduceFn, RearrangeFn, RepeatFn}; // We create a random tensor as input let input = Tensor::randn(&[100, 32, 64], (Kind::Float, Device::Cpu)); // ------------------------------------------------------------------------ // Rearrange operation let output = Rearrange::new("t b c -> b c t")?.apply(&input)?; assert_eq!(output.size(), vec![32, 64, 100]); // Apply rearrange operation directly on the tensor using `RearrangeFn` trait let output = input.rearrange("t b c -> b c t")?; assert_eq!(output.size(), vec![32, 64, 100]); // ------------------------------------------------------------------------ // Perform reduction on first axis let output = Reduce::new("t b c -> b c", Operation::Max)?.apply(&input)?; assert_eq!(output.size(), vec![32, 64]); // Same reduction done directly on the tensor using `ReduceFn` trait let output = input.reduce("t b c -> b c", Operation::Max)?; assert_eq!(output.size(), vec![32, 64]); // ------------------------------------------------------------------------ // We repeat the third axis let output = Repeat::new("t b c -> t b c 3")?.apply(&input)?; assert_eq!(output.size(), vec![100, 32, 64, 3]); // Same as above, but using `RepeatFn` trait and we specify the repeating // axis with a name along with its size and pass it in a slice let output = input.repeat_with_lengths("t b c -> t b c repeat", &[("repeat", 3)])?; assert_eq!(output.size(), vec![100, 32, 64, 3]);
Structs
Rearrange | Reader-friendly reordering of tensors. Includes the functionality of transpose (axes permutation), reshape, squeeze, unsqueeze and other operations |
Reduce | Provides combination of reordering and reduction using reader-friendly notation |
Repeat | Repeat allows reordering elements and repeating them in arbitrary combinations. This includes functionality of repeat, tile, and broadcast functions |
Enums
EinopsError | Different categories of error that can be encountered |
Operation | Specifies the operation used to reduce an axis |
Traits
RearrangeFn | Trait that allows calling the rearrange operation directly on the tensor |
ReduceFn | Trait that allows calling the reduce operation directly on the tensor |
RepeatFn | Trait that allows calling the repeat operation directly on the tensor |