Reading a transformer's forward pass usually means holding shapes in your head. This puts the measured shape on every line, and when it rewrites matmuls and attention as einsum it validates the result before you take it.
A shape-extraction and rewrite pipeline for PyTorch, exposed over MCP. It runs torch.export and ShapeProp across your nn.Module to capture the shape of every intermediate tensor, then annotates the source or rewrites operations as torch.einsum.
- Per-module tensor shapes for every intermediate in the forward pass, not just the inputs
- Your source files back with those shapes inserted as inline comments on the lines they belong to
- Comment mode: shapes rewritten with descriptive dimension names such as batch_size, seq_len and n_embd, plus an explanation of each transformation
- Einsum mode: matmuls and attention operations replaced with torch.einsum, collapsing intermediate reshapes where possible
- A rewrite validated before you accept it — whole file, single module or single function
The repository checked out locally with PyTorch, and a module that torch.export can trace — the shapes come from a real forward pass, not from static analysis. The comment and einsum rewrites run through an LLM, so those two modes need a model behind them.
One command — pip install -e .