Video by Meta Open Source via YouTube
Shape bugs are some of the most painful bugs in ML code, and traditional type systems can’t catch them, because all they tell you is that you have "a tensor."
In this video, Meta engineers Avik Chaudhuri and Steven Troxler introduce static tensor shape checking in Pyrefly: a production-grade static type system that tracks tensor shapes. Starting from just a few annotations at module boundaries, Pyrefly infers the shape of every intermediate value, shows it as an inlay hint in your editor, and flags mismatches before you ever run your code, across PyTorch, JAX, and NumPy code.
TRY IT
• Pyrefly: https://pyrefly.org
• GitHub: https://github.com/facebook/pyrefly
• Docs: getting started with Pyrefly tensor shapes: https://pyrefly.org/en/docs/tensor-shapes/
• Discord: #shape-types channel: https://discord.com/channels/1319086885024567347/1519422996970737885
CHAPTERS
00:00 Intro: static tensor shapes come to Pyrefly
03:16 Demo: a PyTorch cross-attention module with two hidden bugs
06:09 Demo: how Pyrefly uses symbolic arithmetic
07:40 Demo: NumPy array shapes
12:33 Demo: JAX Gaussians and self-attention
13:50: Demo: PyTorch MLP with nn.Sequential
14:51 Demo: PyTorch tracking shapes through a convnet
15:38 Demo: Annotating real-world model codebases using Pyrefly tensor shapes
18:50 Demo: auto-annotating Whisper with the Pyrefly agent skill
21:36 What’s next?
23:11 How to try Pyrefly tensor shape checking today