Skip to content

Community/Ecosystem

Tools for einops users

  • Ein Color — a VS Code extension to color axes in einops/einsum patterns
  • Sonar codechecker provides a rule to statically check einops patterns:

Selected projects implemented with einops

Einops tutorials cover many common usages (cover tutorials first!), but it is also useful to see real projects that apply einops in practice. The projects below illustrate how einops can simplify code in various domains.

  • @lucidrains has a dramatic collection of vision transformers

    • there is a plenty of good examples how to use einops efficiently in your projects

  • lambda networks (non-conventional architecture) implemented by @lucidrains

    • nice demonstration how clearer code can be with einops, even compared to description in the paper
    • implementation and video

  • capsule networks (aka capsnets) implemented in einops

    • this implementation is blazingly fast, concise (3-10 times less code), and memory efficient

  • NuX — normalizing flows in Jax

    • different rearrangement patterns in normalizing flows have nice mapping to einops

  • For video recognition, look at MotionFormer and TimeSFormer implementations

  • For protein folding, see alphafold3-pytorch and implementation of invariant point attention from AlphaFold 2

Community introductions to einops

Tutorial in the AI summer about einops and einsum: https://theaisummer.com/einsum-attention/

Introduction to einops by Kapil Sachdeva https://www.youtube.com/watch?v=xGy75Pjsqzo

Implementing visual transformer in pytorch: https://towardsdatascience.com/implementing-visualttransformer-in-pytorch-184f9f16f632

Refactoring machine learning code, one of posts in a series is devoted to einops: https://www.paepper.com/blog/posts/refactoring-machine-learning-code-einops/

ML TLDR thread on einops: https://twitter.com/mlsummaries/status/1400505282543955970

Book "Deep Reinforcement Learning in Action" by Brandon Brown & Alexander Zai contains an introduction into einops in chapter 10.