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)

Philosophy

TorchShapeFlow is annotation-first and symbolic-first.

• You declare tensor shape contracts with Annotated[torch.Tensor, Shape(...)].
• Symbolic dimensions like "B", "T", and "D" are the default path for config-driven model code.
• Integer dimensions are still useful for fixed semantics like RGB channels or known embedding widths.
• When inference is not possible, the analyzer degrades visibly instead of guessing.

If Pydantic gives structure to data boundaries, TorchShapeFlow aims to do the same for tensor-shape boundaries in deep learning code.

Install

In Claude Code (two commands, no config-file editing):

/plugin marketplace add Davidxswang/torchshapeflow
/plugin install torchshapeflow@torchshapeflow

The first command registers this repo as a plugin marketplace (pulling from main by default). The second installs the torchshapeflow plugin from that marketplace, which wires in an MCP server, an agent skill, and a post-edit hook — your Claude Code then knows how to run tsf check, interpret the structured diagnostics, and propose annotations. No manual .mcp.json editing required.

As a plain Python package (for CLI use or other agent runtimes):

pip install torchshapeflow

Documentation

Full docs at davidxswang.github.io/torchshapeflow

• Quickstart — install and run your first check
• Annotation syntax — how to annotate your tensors
• Supported operators — what is analyzed and what shapes are inferred
• Limitations — what the analyzer does not handle
• For AI coding agents — how Claude Code / Cursor / Copilot / Aider should invoke the CLI and interpret output

Contributing

git clone https://github.com/Davidxswang/torchshapeflow
cd torchshapeflow
make install   # uv sync --extra dev
make check     # format + lint + typecheck + tests

If you want to execute the example PyTorch scripts in examples/, install the separate examples extra:

uv sync --extra dev --extra examples

See docs/development.md for the full development guide: all make targets, CI workflow descriptions, and how to add new operators.