Skip to content

TorchShapeFlow

TorchShapeFlow is a static, AST-based shape analyzer for PyTorch. It reads your Python source, infers tensor shapes from Annotated[..., Shape(...)] contracts, and reports mismatches as structured diagnostics. No execution required.

from typing import Annotated
import torch
from torchshapeflow import Shape

def attention_scores(
    q: Annotated[torch.Tensor, Shape("B", "H", "T", "D")],
    k: Annotated[torch.Tensor, Shape("B", "H", "T", "D")],
) -> Annotated[torch.Tensor, Shape("B", "H", "T", "T")]:
    return q @ k.transpose(-2, -1)

$ tsf check mymodel.py
All clean (1 file checked)

$ tsf check broken.py
broken.py:9:9 error TSF1004 Invalid reshape.

Philosophy

Like Pydantic for data validation, TorchShapeFlow is annotation-first: you declare shape contracts in annotations, and the analyzer verifies consistency. In practice that usually starts with function parameters, then expands into shared shape aliases and annotated local variables as coverage grows. Without annotations, there is nothing to check — and that is by design. You opt in where it matters, starting with forward, and extend coverage incrementally. Symbolic dimensions ("B", "T", "D") are the primary mechanism; the analyzer verifies that operations are consistent without needing concrete sizes. Constants still matter, but mainly for semantically fixed axes like channels, head counts, or embedding widths.

What it does

  • Reads Annotated[torch.Tensor, Shape(...)] contracts from parameters, shape aliases, and annotated local variables
  • Propagates symbolic shapes through supported PyTorch operations
  • Emits diagnostics when shapes are incompatible
  • Provides hover-style shape facts for editor integration

Getting started

For AI coding agents

  • Agent guide — how Claude Code, Cursor, Copilot, and other tool-using LLMs should invoke tsf check / tsf suggest and interpret the output

For contributors

  • Architecture — module map, analysis pipeline, Dim type system
  • Development — make targets, CI, how to add a new operator