PyTorch as a Reference Language: Separating Specification from Implementation

Edward Yang (@ezyang), one of the developers of PyTorch, muses in a note for DevLog on what PyTorch really is in the age of compilers, kernel DSLs and coding AI agents. His starting point is the notion of a reference implementation: a simplified but complete version of a system that sacrifices performance for clarity. A “reference language,” then, is the fabric of APIs and conventions out of which such implementations are cut.

PyTorch is habitually called the lingua franca of deep learning, but, the author notes, there is some confusion here. Reference implementations aren't usually shipped to production — yet people run real training on PyTorch. Kernels are increasingly written in specialized DSLs — so what is the framework's role if it merely glues those kernels together? And if AI sooner or later makes it possible to rewrite any stack from scratch, why does it even matter that the code is written specifically in PyTorch?

One implementation to explore, another to scale

Yang suggests seeing a dual role in PyTorch: both a reference language and an implementation language. As long as the scale is small or the compiler can cope, the reference implementation can go straight to production. But it is ever more natural, he believes, to treat it as a separate artifact against which the correctness of the production implementation is checked. To paraphrase Tolkien: one implementation to explore, another to scale, and one verifier to bind them.

The clearest example is the kernel DSL. Compiler maximalism assumed that the user writes modules in a high-level API and the compiler optimizes everything itself. But for key operations like matrix multiplication and attention, it is hard for a compiler to reach peak performance, and kernel authors spell out tiling and data movement by hand. The high-level API doesn't disappear, though: alongside the optimized kernel they keep a reference implementation in “ordinary” PyTorch and cross-check the results with numerical tests.

Coding agents and the backward pass

The author carries the same logic over to training steps. Autograd is PyTorch's pride, guaranteeing correct derivatives, but at large scale the implicit backward graph becomes a burden: the lion's share of the computation is hidden and can't be touched with the usual debuggers. The new recipe: leave the autograd code as the reference, and use an LLM to generate an explicit forward-backward version that can be optimized separately. Unlike pattern matching, the optimization will never “fall off.” The price for this is a possible divergence between the reference and the production code, so a verifier is needed: bitwise comparison, or structural equivalence of the graphs in the spirit of translation validation.

Yang does not claim the recipe suits everyone: PyTorch as a reference language turned out to be a decent executable specification, but in the end all that matters is the speed of getting the experimental results you need. Yet this lens, in his view, answers Horace He's longstanding question — how to get all the control of eager mode together with the conveniences of graph abstractions.

Source: Lobsters