torchshapeflow 0.2.0__tar.gz

torchshapeflow-0.2.0/.github/workflows/build-artifacts.yml

@@ -0,0 +1,41 @@
+name: Build Artifacts
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  python-package:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+      - run: uv build
+      - uses: actions/upload-artifact@v7
+        with:
+          name: python-dist
+          path: dist/*
+
+  vscode-extension:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-node@v6
+        with:
+          node-version: "24"
+          cache: "npm"
+          cache-dependency-path: extensions/vscode/package-lock.json
+      - working-directory: extensions/vscode
+        run: npm ci
+      - working-directory: extensions/vscode
+        run: npm run package
+      - uses: actions/upload-artifact@v7
+        with:
+          name: vscode-extension
+          path: extensions/vscode/dist/*.vsix

torchshapeflow-0.2.0/.github/workflows/ci.yml

@@ -0,0 +1,43 @@
+name: CI
+
+on:
+  push:
+  pull_request:
+
+jobs:
+  check:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+      - uses: actions/setup-node@v6
+        with:
+          node-version: "24"
+          cache: "npm"
+          cache-dependency-path: extensions/vscode/package-lock.json
+      - run: uv sync --extra dev
+      - run: uv run ruff format . --check
+      - run: uv run ruff check .
+      - run: uv run mypy .
+      - run: uv run mkdocs build
+      - working-directory: extensions/vscode
+        run: npm ci
+      - working-directory: extensions/vscode
+        run: npm run build
+
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+      - uses: actions/setup-python@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+      - run: uv sync --extra dev
+      - run: uv run pytest -q

torchshapeflow-0.2.0/.github/workflows/docs.yml

@@ -0,0 +1,42 @@
+name: Docs
+
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+
+permissions:
+  contents: read
+  pages: write
+  id-token: write
+
+concurrency:
+  group: pages
+  cancel-in-progress: true
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+      - run: uv sync --extra dev
+      - run: uv run mkdocs build
+      - uses: actions/upload-pages-artifact@v4
+        with:
+          path: site
+
+  deploy:
+    if: github.ref == 'refs/heads/main'
+    needs: build
+    runs-on: ubuntu-latest
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    steps:
+      - id: deployment
+        uses: actions/deploy-pages@v4

torchshapeflow-0.2.0/.github/workflows/release.yml

@@ -0,0 +1,144 @@
+name: Release
+
+on:
+  push:
+    tags:
+      - "v*"
+  workflow_dispatch:
+
+jobs:
+  check-secrets:
+    runs-on: ubuntu-latest
+    outputs:
+      has-vsce-pat: ${{ steps.check.outputs.has-vsce-pat }}
+      has-ovsx-pat: ${{ steps.check.outputs.has-ovsx-pat }}
+    steps:
+      - id: check
+        env:
+          VSCE_PAT: ${{ secrets.VSCE_PAT }}
+          OVSX_PAT: ${{ secrets.OVSX_PAT }}
+        run: |
+          echo "has-vsce-pat=${{ secrets.VSCE_PAT != '' }}" >> "$GITHUB_OUTPUT"
+          echo "has-ovsx-pat=${{ secrets.OVSX_PAT != '' }}" >> "$GITHUB_OUTPUT"
+
+  build-python:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: astral-sh/setup-uv@v7
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+      - run: uv build
+      - uses: actions/upload-artifact@v7
+        with:
+          name: python-dist
+          path: dist/*
+
+  build-extension:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-node@v6
+        with:
+          node-version: "24"
+          cache: "npm"
+          cache-dependency-path: extensions/vscode/package-lock.json
+      - working-directory: extensions/vscode
+        run: npm ci
+      - working-directory: extensions/vscode
+        run: npm run package
+      - uses: actions/upload-artifact@v7
+        with:
+          name: vscode-extension
+          path: extensions/vscode/dist/*.vsix
+
+  publish-vscode-marketplace:
+    if: needs.check-secrets.outputs.has-vsce-pat == 'true' && !contains(github.ref, '-rc') && !contains(github.ref, '-test')
+    needs: [check-secrets, build-extension]
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-node@v6
+        with:
+          node-version: "24"
+          cache: "npm"
+          cache-dependency-path: extensions/vscode/package-lock.json
+      - working-directory: extensions/vscode
+        run: npm ci
+      - working-directory: extensions/vscode
+        env:
+          VSCE_PAT: ${{ secrets.VSCE_PAT }}
+        run: npx @vscode/vsce publish -p "$VSCE_PAT"
+
+  publish-open-vsx:
+    if: needs.check-secrets.outputs.has-ovsx-pat == 'true' && !contains(github.ref, '-rc') && !contains(github.ref, '-test')
+    needs: [check-secrets, build-extension]
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-node@v6
+        with:
+          node-version: "24"
+          cache: "npm"
+          cache-dependency-path: extensions/vscode/package-lock.json
+      - working-directory: extensions/vscode
+        run: npm ci
+      - working-directory: extensions/vscode
+        run: npm run package
+      - working-directory: extensions/vscode
+        env:
+          OVSX_PAT: ${{ secrets.OVSX_PAT }}
+        run: npx ovsx publish dist/torchshapeflow.vsix -p "$OVSX_PAT"
+
+  publish-testpypi:
+    if: contains(github.ref, '-rc') || contains(github.ref, '-test')
+    needs: build-python
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write
+    environment:
+      name: testpypi
+    steps:
+      - uses: actions/download-artifact@v8
+        with:
+          name: python-dist
+          path: dist
+      - uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          repository-url: https://test.pypi.org/legacy/
+
+  publish-pypi:
+    if: "!contains(github.ref, '-rc') && !contains(github.ref, '-test')"
+    needs: build-python
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write
+    environment:
+      name: pypi
+    steps:
+      - uses: actions/download-artifact@v8
+        with:
+          name: python-dist
+          path: dist
+      - uses: pypa/gh-action-pypi-publish@release/v1
+
+  github-release:
+    needs: [build-python, build-extension]
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    steps:
+      - uses: actions/download-artifact@v8
+        with:
+          name: python-dist
+          path: dist
+      - uses: actions/download-artifact@v8
+        with:
+          name: vscode-extension
+          path: extensions/vscode/dist
+      - uses: softprops/action-gh-release@v2
+        with:
+          files: |
+            dist/*
+            extensions/vscode/dist/*.vsix

torchshapeflow-0.2.0/.gitignore

@@ -0,0 +1,13 @@
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.venv/
+build/
+dist/
+htmlcov/
+site/
+*.pyc
+__pycache__/
+node_modules/
+extensions/vscode/out/
+*.vsix

torchshapeflow-0.2.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Xuesong Wang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torchshapeflow-0.2.0/Makefile

@@ -0,0 +1,52 @@
+PYTHON ?= python
+
+.PHONY: install format lint typecheck test check docs docs-serve build python-dist extension-build extension-package bump-patch bump-minor bump-major clean
+
+install:
+	uv sync --extra dev
+
+format:
+	uv run ruff format .
+
+lint:
+	uv run ruff check . --fix
+
+typecheck:
+	uv run mypy .
+
+test:
+	uv run pytest -q
+
+check: format lint typecheck test
+
+docs:
+	uv run mkdocs build
+
+docs-serve:
+	uv run mkdocs serve
+
+build: python-dist extension-package
+
+python-dist:
+	uv build
+
+extension-build:
+	cd extensions/vscode && npm ci && npm run build
+
+extension-package:
+	cd extensions/vscode && npm ci && npm run package
+
+bump-patch:
+	uv run python scripts/bump_version.py patch
+	uv lock
+
+bump-minor:
+	uv run python scripts/bump_version.py minor
+	uv lock
+
+bump-major:
+	uv run python scripts/bump_version.py major
+	uv lock
+
+clean:
+	rm -rf .pytest_cache .mypy_cache .ruff_cache build dist site extensions/vscode/dist extensions/vscode/out

torchshapeflow-0.2.0/PKG-INFO

@@ -0,0 +1,113 @@
+Metadata-Version: 2.4
+Name: torchshapeflow
+Version: 0.2.0
+Summary: Static AST-based PyTorch tensor shape analysis.
+Project-URL: Homepage, https://github.com/Davidxswang/torchshapeflow
+Project-URL: Repository, https://github.com/Davidxswang/torchshapeflow
+Project-URL: Issues, https://github.com/Davidxswang/torchshapeflow/issues
+Project-URL: Documentation, https://davidxswang.github.io/torchshapeflow/
+Author: TorchShapeFlow Contributors
+License: MIT License
+        
+        Copyright (c) 2026 Xuesong Wang
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Requires-Python: >=3.10
+Requires-Dist: typer>=0.16.0
+Provides-Extra: dev
+Requires-Dist: mkdocs-material>=9.5.0; extra == 'dev'
+Requires-Dist: mkdocs>=1.6.0; extra == 'dev'
+Requires-Dist: mypy>=1.18.0; extra == 'dev'
+Requires-Dist: pytest>=8.3.0; extra == 'dev'
+Requires-Dist: ruff>=0.11.0; extra == 'dev'
+Requires-Dist: torch>=2.10.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# TorchShapeFlow
+
+[![CI](https://github.com/Davidxswang/torchshapeflow/actions/workflows/ci.yml/badge.svg)](https://github.com/Davidxswang/torchshapeflow/actions/workflows/ci.yml)
+[![PyPI](https://img.shields.io/pypi/v/torchshapeflow)](https://pypi.org/project/torchshapeflow/)
+[![Python](https://img.shields.io/pypi/pyversions/torchshapeflow)](https://pypi.org/project/torchshapeflow/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
+
+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.
+
+```python
+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]
+```
+
+```bash
+$ tsf check mymodel.py
+mymodel.py: ok
+```
+
+## Install
+
+```bash
+pip install torchshapeflow
+```
+
+## Documentation
+
+Full docs at **[davidxswang.github.io/torchshapeflow](https://davidxswang.github.io/torchshapeflow)**
+
+- [Quickstart](https://davidxswang.github.io/torchshapeflow/quickstart/) — install and run your first check
+- [Annotation syntax](https://davidxswang.github.io/torchshapeflow/syntax/) — how to annotate your tensors
+- [Supported operators](https://davidxswang.github.io/torchshapeflow/operators/) — what is analyzed and what shapes are inferred
+- [Limitations](https://davidxswang.github.io/torchshapeflow/limitations/) — what the analyzer does not handle
+
+## Contributing
+
+```bash
+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](docs/development.md) for the full development guide: all make targets, CI workflow descriptions, and how to add new operators.
+
+## Release
+
+See [RELEASING.md](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

torchshapeflow-0.2.0/README.md

@@ -0,0 +1,71 @@
+# TorchShapeFlow
+
+[![CI](https://github.com/Davidxswang/torchshapeflow/actions/workflows/ci.yml/badge.svg)](https://github.com/Davidxswang/torchshapeflow/actions/workflows/ci.yml)
+[![PyPI](https://img.shields.io/pypi/v/torchshapeflow)](https://pypi.org/project/torchshapeflow/)
+[![Python](https://img.shields.io/pypi/pyversions/torchshapeflow)](https://pypi.org/project/torchshapeflow/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
+
+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.
+
+```python
+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]
+```
+
+```bash
+$ tsf check mymodel.py
+mymodel.py: ok
+```
+
+## Install
+
+```bash
+pip install torchshapeflow
+```
+
+## Documentation
+
+Full docs at **[davidxswang.github.io/torchshapeflow](https://davidxswang.github.io/torchshapeflow)**
+
+- [Quickstart](https://davidxswang.github.io/torchshapeflow/quickstart/) — install and run your first check
+- [Annotation syntax](https://davidxswang.github.io/torchshapeflow/syntax/) — how to annotate your tensors
+- [Supported operators](https://davidxswang.github.io/torchshapeflow/operators/) — what is analyzed and what shapes are inferred
+- [Limitations](https://davidxswang.github.io/torchshapeflow/limitations/) — what the analyzer does not handle
+
+## Contributing
+
+```bash
+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](docs/development.md) for the full development guide: all make targets, CI workflow descriptions, and how to add new operators.
+
+## Release
+
+See [RELEASING.md](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

torchshapeflow-0.2.0/docs/architecture.md

@@ -0,0 +1,85 @@
+# Architecture
+
+## Analysis pipeline
+
+TorchShapeFlow analyzes Python source in a single pass per file:
+
+1. **Parse** — `ast.parse` converts source text into an AST module (`parser.parse_source`).
+2. **Collect module specs** — `_collect_class_specs` walks class `__init__` bodies to find `nn.Linear`, `nn.Conv2d`, `nn.Embedding`, `nn.MaxPool2d`, `nn.AvgPool2d`, `nn.Sequential`, `nn.MultiheadAttention`, and passthrough module assignments, recording their constructor arguments as spec values.
+3. **Seed shape environment** — for each function (or `forward` method), annotated parameters are parsed via `parser.parse_tensor_annotation` and added to the environment `env: dict[str, Value]`.
+4. **Propagate shapes** — `_analyze_statement` walks the function body statement by statement. For each assignment, `_eval_expr` evaluates the right-hand side, dispatching to the appropriate rule function. Results are stored back into `env`.
+5. **Emit results** — diagnostics and hover facts accumulate in a `ModuleContext` and are returned as a `FileReport`.
+
+## Module map
+
+| Module | Responsibility |
+|---|---|
+| `model.py` | All core data types (`Dim` variants, `TensorShape`, `TensorValue`, `TensorTupleValue`, `LinearSpec`, `Conv2dSpec`, `PassthroughSpec`, `EmbeddingSpec`, `Pool2dSpec`, `SequentialSpec`, `MultiheadAttentionSpec`, `ModuleSpec`, `Value`). Shape arithmetic: `product_dim`, `quotient_dim`, `sum_dim`, `broadcast_shapes`, `batch_matmul_shape`, `normalize_index`. |
+| `annotations.py` | Public `Shape` class used in `Annotated[Tensor, Shape(...)]`. |
+| `parser.py` | Parses `Annotated[Tensor, Shape(...)]` annotation AST nodes into `TensorValue`. Raises `AnnotationParseError` on malformed annotations. |
+| `analyzer.py` | Main AST walker. Manages the shape environment, dispatches to rule functions, emits diagnostics via `ModuleContext`. |
+| `diagnostics.py` | `Diagnostic` dataclass and `Severity` type alias (`"error" \| "warning"`). |
+| `report.py` | `FileReport` (list of diagnostics + hover facts per file) and `HoverFact` (inferred shape at a source location). |
+| `cli.py` | Typer CLI. `tsf check` runs the analyzer and formats output. `tsf version` prints the package version. |
+| `rules/__init__.py` | Re-exports all public inference functions. |
+| `rules/shape_ops.py` | `infer_permute`, `infer_transpose`, `infer_reshape`, `infer_flatten`, `infer_squeeze`, `infer_unsqueeze`, `infer_size`, `infer_cat`, `infer_stack`, `infer_matmul`, `infer_mm`, `infer_movedim`, `infer_reduction`, `infer_chunk`, `infer_split`, `infer_einsum`. |
+| `rules/broadcasting.py` | `infer_binary_broadcast` — wraps `broadcast_shapes` for element-wise ops. |
+| `rules/linear.py` | `infer_linear` for `nn.Linear`. |
+| `rules/conv2d.py` | `infer_conv2d` for `nn.Conv2d`. |
+| `rules/embedding.py` | `infer_embedding` for `nn.Embedding`. |
+| `rules/pool2d.py` | `infer_pool2d` for `nn.MaxPool2d` and `nn.AvgPool2d`. |
+| `rules/indexing.py` | `infer_subscript` for tensor subscript and shape-tuple indexing. |
+| `rules/common.py` | Shared AST helpers: `int_from_ast`, `qualified_name`, `dim_from_value`, `tuple_index`, `spatial_output_dim`. |
+| `utils/paths.py` | `collect_python_files` — recursive `.py` file discovery. |
+
+## Dim type hierarchy
+
+```
+Dim  (TypeAlias)
+├── ConstantDim(value: int)       — a fixed integer size, e.g. 32
+├── SymbolicDim(name: str)        — a named unknown size, e.g. "B"
+├── ExpressionDim(expr: str)      — a derived expression, e.g. "4*B" or "(B*C)/4"
+└── UnknownDim(token: str)        — explicitly unresolvable
+```
+
+Shape arithmetic returns `ConstantDim` when all operands are constant and `ExpressionDim` otherwise. Expressions are stored as strings and compared structurally.
+
+## Shape environment
+
+The environment `env: dict[str, Value]` maps variable names to their inferred `Value`:
+
+```
+Value  (TypeAlias)
+├── TensorValue(shape: TensorShape, origin: str | None)
+├── ShapeTupleValue(dims: tuple[Dim, ...])   — result of x.shape or x.size()
+├── IntegerValue(value: int | None)           — result of x.ndim or x.size(i)
+├── TensorTupleValue(tensors: tuple[TensorValue, ...])  — result of chunk/split/MHA
+│
+│   ModuleSpec  (TypeAlias — stored in module_specs and env)
+├── LinearSpec(in_features, out_features)     — nn.Linear
+├── Conv2dSpec(in_channels, out_channels, kernel_size, stride, padding, dilation)  — nn.Conv2d
+├── PassthroughSpec()                         — shape-preserving modules (BatchNorm, ReLU, …)
+├── EmbeddingSpec(embedding_dim)              — nn.Embedding
+├── Pool2dSpec(kernel_size, stride, padding, dilation)  — nn.MaxPool2d / nn.AvgPool2d
+├── SequentialSpec(specs: tuple[ModuleSpec, ...])       — nn.Sequential
+└── MultiheadAttentionSpec(embed_dim, num_heads, batch_first)  — nn.MultiheadAttention
+```
+
+Spec values are stored in `module_specs` (keyed by attribute name) when their constructor is parsed from `__init__`. When `self.linear(x)` is called, the analyzer looks up `"linear"` in `module_specs`, retrieves the spec, and calls the appropriate inference function. Module aliases (`m = self.linear; m(x)`) are also supported: spec values stored in `env` are looked up before falling through to `func_sigs`.
+
+## Diagnostic codes
+
+| Code | Trigger |
+|---|---|
+| `TSF1001` | Annotation parse error (malformed `Annotated` or `Shape`) |
+| `TSF1003` | Incompatible `matmul` / `bmm` shapes |
+| `TSF1004` | Invalid `reshape` or `flatten` dimensions |
+| `TSF1005` | Invalid `cat` or `stack` dimensions or mismatched shapes |
+| `TSF1006` | Broadcasting incompatibility |
+| `TSF1007` | `nn.Linear`, `nn.Conv2d`, or `nn.MaxPool2d`/`AvgPool2d` input shape mismatch |
+| `TSF1008` | Invalid `permute`, `transpose`, `squeeze`, or `unsqueeze` dimensions |
+| `TSF1009` | Return shape does not match the declared return type annotation |
+
+## Adding a new operator
+
+See [Development → Adding a new operator](development.md#adding-a-new-operator).