Einops.jl

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.

Description	Einops	Base primitives
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 singleton	@rearrange(x, "1 ... -> ...")	dropdims(x, dims=1)
Funky repeat	@repeat(x, "... -> 2 ... 3")	repeat( reshape(x, 1, size(x)...), 2, ntuple(Returns(1), ndims(x))..., 3)
Setup for Multi-Head Attention	@rearrange(q, "(d h) l b -> d l (h b)"; d=head_dim)	reshape( permutedims( reshape(q, head_dim, :, size(q)[2:3]...), (1, 3, 2, 4)), head_dim, size(q, 2), :)
Setup for Grouped-Query Attention	@repeat(k, "(d h) l b -> d l (r h b)"; d=head_dim, r=repeats)	reshape( repeat( permutedims( reshape(k, head_dim, :, size(k)[2:3]...), (1, 3, 2, 4)), inner=(1, 1, repeats, 1)), head_dim, size(k, 2), :)

Contributing

Contributions are welcome! Please feel free to open an issue to report a bug or start a discussion.