[ANN] Einops.jl: Concise, declarative tensor operations

Package Announcements
, , ,
1

I’m excited to share that I’ve been working on Einops.jl, bringing the popular einops tensor manipulation syntax to Julia. While Julia’s array interface is already excellent thanks to multiple dispatch, there’s still value in having a more declarative way to express complex tensor operations.

What it does

Instead of chaining primitives:

reshape(repeat(permutedims(reshape(x, head_dim, :, size(x)[2:3]...), (1, 3, 2, 4)), inner=(1, 1, repeats, 1)), head_dim, size(x, 2), :)

You declare the pattern:

repeat(x, einops"(dim head) len batch -> dim len (r head batch)", d=head_dim, r=repeats)

# or without `@einops_str`:

repeat(x, ((:dim, :head), :len, :batch) --> (:dim, :len, (:r, :head, :batch)), d=head_dim, r=repeats)

(I’m slightly more generous to the Base primitives in the README)

Key points

  • Zero overhead: Expands directly to Base primitives like reshape, permutedims, repeat using generated functions, and thus it should necessarily be compatible with all array types.

  • No unnecessary sub-operations: Only applies operations when necessary (e.g. don’t permute if just reshaping).

  • Type stable: Patterns are embedded in a ArrowPattern{L,R} or Val{pattern} for specialization, depending on the function, but the einops string macro is meant to hide this.

  • Self-documenting: Dimensions are explicitly named.

  • Familiar syntax: If you’ve used einops in Python, you already know the API, and porting to Julia becomes a breeze. Even if you haven’t, it will quickly click.

The package handles common operations like rearranging, reducing, and repeating dimensions (main focus) but also some extra stuff for parity, including a binding to the neat OMEinsum.jl package.

There is also TensorCast.jl, a superset of Einops.jl in terms of functionality, which uses a more sophisticated explicit macro syntax. In my experience, macros often slightly obscure what’s going on, so I created Einops.jl to serve a different niche – predictable and lightweight, with a well-defined scope.

Still exploring lazy permutations, but the core functionality is solid and ready for use. Feedback and contributions are welcome!

Docs | GitHub | einops

9 Likes
2

Welcome Anton!

Just a note here, you’re also using a macro, it just happens to be a string macro.

1 Like
3

Yes, as I also noted. Perhaps I should’ve mentioned that there’s an alternative syntax using a --> operator, which the einops string macro expands to. You can write e.g. (:a, :b, ..) --> (:b, :a, ..) (hat tip, EllipsisNotation.jl) which puts both sides in type parameters so they can be specialized upon. This syntax is also just made up, of course, but I’ve come to favor the string macro simply because it’s slightly easier to write because of the colons, commas, and superfluous parentheses in the --> version. See also CallMode.jl if you despise parentheses.

4

Awesome! Are you optimizing the contraction schedule via OMEinsumContractionOrders.jl or EinExprs.jl or something similar?

1 Like
5

Thanks! and no, not currently, but I just opened a PR: Optimize einsum contraction order by AntonOresten · Pull Request #30 · MurrellGroup/Einops.jl · GitHub.
Any help on this front would be appreciated. Contraction order is nuanced, especially when I don’t want it to be user-facing. I acknowledge that Julia not in need of one more “One More Einsum” package! :wink: