torchshapeflow 0.2.1

Static AST-based PyTorch tensor shape analysis.

TorchShapeFlow

TorchShapeFlow is a static, AST-based shape analyzer for PyTorch. It reads your Python source — no execution required — infers tensor shapes through your code, and reports mismatches as structured diagnostics.

from typing import Annotated
import torch
import torch.nn as nn
from torchshapeflow import Shape

class Net(nn.Module):
    def __init__(self):
        self.conv = nn.Conv2d(3, 8, 3, padding=1)
        self.linear = nn.Linear(8 * 32 * 32, 10)

    def forward(self, x: Annotated[torch.Tensor, Shape("B", 3, 32, 32)]):
        y = self.conv(x)      # inferred: [B, 8, 32, 32]
        z = y.flatten(1)      # inferred: [B, 8192]
        return self.linear(z) # inferred: [B, 10]

$ tsf check mymodel.py
mymodel.py: ok

Install

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

Contributing

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

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

Release

See RELEASING.md for the full release procedure.

Build commands:

• make python-dist — wheel and sdist into dist/
• make extension-package — VS Code extension .vsix
• make build — both

Marketplace publishing in the release workflow is gated on GitHub Actions secrets:

• VSCE_PAT for the VS Code Marketplace
• OVSX_PAT for Open VSX