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), :)

Axis lengths

A keyword gives a name its length. An integer is an axis known only by its length: 1 is an axis of length one (inside parentheses it adds nothing), and any other integer can be reduced by @reduce or created by @repeat. $ interpolates such a length, as if the value were written there, while keeping it out of the pattern's type, so it can be computed at run time:

@rearrange(q, "(d h) l b -> d l h b"; d=head_dim)
@reduce(mean, x, "(h $k) (w $k) c -> h w c")  # k×k average pooling

Every length a pattern fixes is checked against the input:

julia> @reduce(sum, rand(3, 2), "4 b -> b")
ERROR: DimensionMismatch: Input axis 1 has length 3, but the pattern expects 4

Contributing

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