Talk: Tensor Shapes in the Type System | Pyrefly

Skip to main content

Why aren't tensor shapes part of the type system?

When you write a PyTorch model, the hardest part of composing tensor ops is keeping track of their shapes. The standard practice today is to write shapes down as comments as you go. This talk introduces an experimental Pyrefly feature that brings tensor shapes into Python's type system, so those shape annotations can become inferred type hints instead of comments.

The talk was originally presented at the PyCon US 2026 Typing Summit. The slides and edited transcript are provided below for your convenience.

📎 Slides (PDF)

Want to give it a spin? This feature is available in Pyrefly today, and we'd love to hear your feedback and suggestions.

The problem​

Why are we working on this?

If you squint, a PyTorch model is just a function from a tensor to another tensor. In between is a whole bunch of transformations: you apply so-called tensor ops, and all they're doing is transforming these tensors from your input to your output. So you can think of it as a straight-line program, or as a graph.

The hardest part of model development—unless you're a researcher trying to change the architecture—is that once you know the architecture and you're just trying to compose these ops, you have to think about their shapes. Just like in regular programming, where you think about what function to call and how to compose their types, shapes are the unit of types in PyTorch programs, or in models in general.

The standard practice, which holds today, is that you write a bunch of comments in your code to track the shapes, because it's really difficult. These ops are non-trivial—they're doing complicated transformations—so you write down the shapes as comments as you go.

Here is a snippet of a thing called nanoGPT, written by Andrej Karpathy. Some of you might know him; he's a really famous AI researcher, but he's also a very good educator. Very early on, when GPT-2 came out, he came up with this tiny repository called nanoGPT, where he goes through the mechanics of how to put together a real model—but in a very educational way, so people can understand how the different parts of a classic LLM come together.

Even if you don't follow what's going on line by line, note the green comments. He's following a very standard convention of explaining shapes: when he creates a tensor and assigns it to pos, it has shape T. Then he puts it through a bunch of transformer calls, goes into a loop, and calls a bunch of blocks. These blocks are modules that themselves have other tensor ops inside them. Throughout, he's writing what look like tuples—B, T, embedding—and these are all shapes of tensors at that line.

So the question we're trying to answer is: wouldn't it be nice if, just like regular types, you had type hints provided by the type checker, so you don't have to write comments anymore?

This is the same snippet of code—and I'll show more snippets here and elsewhere—where the types just get inferred for you. This is not magic. You are annotating the boundaries of your functions, but that's common practice. We use generics—and I'll talk about what kind of generics we use—but once you do that, the types of all the locals get inferred, and the shapes all line up as inline type hints.

How it works​

That goes into the why. Let me talk at a high level about how we did this. There are three main ideas. I won't get into a lot of detail, but listen to these three ideas—you could have pieced this together in your heads as well.

Idea 1: Symbolic arithmetic at the type level​

The first idea has nothing to do with tensors per se. It introduces symbolic arithmetic at the type system level. We have a new kind of type argument that represents dimension sizes. Generic types take type arguments; this time, the way to think about this type argument is that it's an integer—but it could be an arbitrary integer, and therefore it could be a symbol, and the symbol is quantified as a type argument somewhere. So it could be a normal integer, it could be a symbol, or it could be an arithmetic expression on these symbols. These are all kinds of int, but they're symbolic expressions.

One key design decision I took pretty early is that I didn't want to throw a SAT solver into the type checker. There is no point where we ask a question like "is there an X such that 13 = X + 6?" Instead, we take the view that these could be expressions. We do need to decide equality between symbolic expressions, and the way we do that is by a normalization procedure: we simplify the two sides and then check for syntactic equality. The normalization procedure knows a bunch of equations that should hold, and it does its best. There's no claim that all of this is decidable, but it works pretty well in practice, and that's the limit to which we go. We do not introduce so-called existential types in this system.

Idea 2: Two user-facing types​

The second...