Einops.jl brings the readable and concise tensor operations of einops to Julia, expanding to primitives like reshape, permutedims, and repeat, while avoiding double-wrappers when possible using Rewrap.jl and remaining type stable.

Quickstart

using Pkg
Pkg.add("Einops")
using Einops

x = rand(2, 3, 4)
y = @rearrange(x, "a b c -> (a b) c")

Einops vs Base primitives

The snippets below show identical transformations expressed first with a single Einops expression and then with “hand-rolled” Julia primitives. Notice how Einops collapses multiple e.g. reshape / permutedims / dropdims / repeat calls into a single, declarative statement, while still expanding to such primitives under the hood and avoiding double-wrappers and no-ops.

Flatten first two dimensions

@rearrange(x, "a b c -> (a b) c")

reshape(x, :, size(x, 3))

Permute first two dimensions

@rearrange(x, "a b c -> b a c")

permutedims(x, (2, 1, 3))

Permute and flatten

@rearrange(x, "a b -> (b a)")

vec(permutedims(x))

Remove first dimension…