Computing sharding with einsum : ezyang's blog
Mental arithmetic in grade school (e.g., memorizing your times tables) is typically justified on the grounds that facility in basic calculations makes it easier to focus on higher-level problems that require being able to do these manipulations. When working on DTensor, I have also found it important to be able to quickly calculate what shardings you get when you do matrix multiplies on sharded tensors. Without being able to do this quickly and accurately, working through examples becomes a slog. I’ve also found that while diagrammatic approaches (e.g., drawing a matrix and slicing it into shards) are intuitive, they are slow and unwieldy to do calculations with. Recently, I’ve found that working on sharding with einsum is nice and efficient, and I hope to persuade you to do it this way when you need to reason about sharding! This post somewhat overlaps with Sharded Matrices and How to Multiply Them, but with some different emphasis and some different notation. Einstein summation is a
Computing sharding with einsum January 25, 2026 Mental arithmetic in grade school (e.g., memorizing your times tables) is typically justified on the grounds that facility in basic calculations makes it easier to focus on higher-level problems that require being able to do these manipulations. When working on DTensor, I have also found it important to be able to quickly calculate what shardings you get when you do matrix multiplies on sharded tensors. Without being able to do this quickly and accurately, working through examples becomes a slog. I've also found that while diagrammatic approaches
Explore this link on the map →saved by
related reading
- Einsum is easy and usefulejenner.com
- Inside NVIDIA GPUs: Anatomy of high performance matmul kernels - Aleksa Gordićaleksagordic.com
- Tim Rocktäschelrockt.ai
- Tim Rocktäschelrockt.github.io
- irhum.github.io - Tensor Parallelism with jax.pjitirhum.github.io
- Tiny TPUtinytpu.com
- PyTorch internals : ezyang's blogblog.ezyang.com
- Transformer Math 101 | EleutherAI Blogblog.eleuther.ai
- 2404.17625arxiv.org
- Chapter 0: Fundamentals - ARENAlearn.arena.education
- From Pairwise to Higher Order Tensor Operations on GPUs – Springtail Blogspringtail.ai
- An Intuitive Guide to Linear Algebra – BetterExplainedbetterexplained.com