torchtyc 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- torchtyc-0.1.0/.github/workflows/ci.yml +48 -0
- torchtyc-0.1.0/.gitignore +11 -0
- torchtyc-0.1.0/LICENSE +21 -0
- torchtyc-0.1.0/PKG-INFO +305 -0
- torchtyc-0.1.0/README.md +282 -0
- torchtyc-0.1.0/bench/RESULTS.md +95 -0
- torchtyc-0.1.0/bench/bench.py +249 -0
- torchtyc-0.1.0/flake.lock +27 -0
- torchtyc-0.1.0/flake.nix +50 -0
- torchtyc-0.1.0/pyproject.toml +49 -0
- torchtyc-0.1.0/src/torchtyc/__init__.py +20 -0
- torchtyc-0.1.0/src/torchtyc/annotations.py +215 -0
- torchtyc-0.1.0/src/torchtyc/binding.py +479 -0
- torchtyc-0.1.0/src/torchtyc/cli.py +180 -0
- torchtyc-0.1.0/src/torchtyc/config.py +134 -0
- torchtyc-0.1.0/src/torchtyc/diagnostics.py +150 -0
- torchtyc-0.1.0/src/torchtyc/discovery.py +728 -0
- torchtyc-0.1.0/src/torchtyc/einops_rules.py +165 -0
- torchtyc-0.1.0/src/torchtyc/engine.py +409 -0
- torchtyc-0.1.0/src/torchtyc/formats.py +164 -0
- torchtyc-0.1.0/src/torchtyc/lsp.py +501 -0
- torchtyc-0.1.0/src/torchtyc/tracing.py +706 -0
- torchtyc-0.1.0/src/torchtyc/worker.py +391 -0
- torchtyc-0.1.0/tests/fixtures/both.py +7 -0
- torchtyc-0.1.0/tests/fixtures/layers.py +51 -0
- torchtyc-0.1.0/tests/test_annotations.py +71 -0
- torchtyc-0.1.0/tests/test_binding.py +354 -0
- torchtyc-0.1.0/tests/test_config.py +82 -0
- torchtyc-0.1.0/tests/test_discovery.py +346 -0
- torchtyc-0.1.0/tests/test_einops_rules.py +114 -0
- torchtyc-0.1.0/tests/test_engine.py +1296 -0
- torchtyc-0.1.0/tests/test_formats.py +58 -0
- torchtyc-0.1.0/tests/test_lint.py +153 -0
- torchtyc-0.1.0/tests/test_lsp.py +438 -0
|
@@ -0,0 +1,48 @@
|
|
|
1
|
+
name: CI
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
branches: [main]
|
|
6
|
+
pull_request:
|
|
7
|
+
|
|
8
|
+
jobs:
|
|
9
|
+
test:
|
|
10
|
+
runs-on: ubuntu-latest
|
|
11
|
+
strategy:
|
|
12
|
+
fail-fast: false
|
|
13
|
+
matrix:
|
|
14
|
+
python: ["3.11", "3.12", "3.13"]
|
|
15
|
+
env:
|
|
16
|
+
UV_PYTHON: ${{ matrix.python }}
|
|
17
|
+
steps:
|
|
18
|
+
- uses: actions/checkout@v4
|
|
19
|
+
- uses: astral-sh/setup-uv@v5
|
|
20
|
+
with:
|
|
21
|
+
enable-cache: true
|
|
22
|
+
- run: uv python install ${{ matrix.python }}
|
|
23
|
+
- run: uv sync --all-extras --dev
|
|
24
|
+
- run: uv run ruff check src tests bench
|
|
25
|
+
- run: uv run pytest -q
|
|
26
|
+
|
|
27
|
+
self-check:
|
|
28
|
+
# torchtyc checking its own fixtures, which is also the smoke test for the
|
|
29
|
+
# github output format.
|
|
30
|
+
runs-on: ubuntu-latest
|
|
31
|
+
steps:
|
|
32
|
+
- uses: actions/checkout@v4
|
|
33
|
+
- uses: astral-sh/setup-uv@v5
|
|
34
|
+
- run: uv sync --all-extras --dev
|
|
35
|
+
# The fixtures are deliberately wrong, so exit 1 (findings) is the
|
|
36
|
+
# expected result. Anything else is torchtyc itself failing.
|
|
37
|
+
- run: |
|
|
38
|
+
status=0
|
|
39
|
+
uv run torchtyc check tests/fixtures --format github > annotations.txt || status=$?
|
|
40
|
+
cat annotations.txt
|
|
41
|
+
if [ "$status" -gt 1 ]; then
|
|
42
|
+
echo "torchtyc exited $status, which is a torchtyc failure" >&2
|
|
43
|
+
exit "$status"
|
|
44
|
+
fi
|
|
45
|
+
if ! grep -q "^::error " annotations.txt; then
|
|
46
|
+
echo "the fixtures produced no github annotation" >&2
|
|
47
|
+
exit 1
|
|
48
|
+
fi
|
torchtyc-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Bhaswata
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
torchtyc-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,305 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: torchtyc
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Shape checking for PyTorch, from your jaxtyping annotations
|
|
5
|
+
Project-URL: Homepage, https://github.com/bhaswata08/torchtyc
|
|
6
|
+
Project-URL: Repository, https://github.com/bhaswata08/torchtyc
|
|
7
|
+
Project-URL: Issues, https://github.com/bhaswata08/torchtyc/issues
|
|
8
|
+
License: MIT
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Requires-Python: >=3.11
|
|
11
|
+
Requires-Dist: jaxtyping>=0.3
|
|
12
|
+
Provides-Extra: all
|
|
13
|
+
Requires-Dist: einops>=0.8; extra == 'all'
|
|
14
|
+
Requires-Dist: pygls>=1.3; extra == 'all'
|
|
15
|
+
Requires-Dist: watchfiles>=0.24; extra == 'all'
|
|
16
|
+
Provides-Extra: einops
|
|
17
|
+
Requires-Dist: einops>=0.8; extra == 'einops'
|
|
18
|
+
Provides-Extra: lsp
|
|
19
|
+
Requires-Dist: pygls>=1.3; extra == 'lsp'
|
|
20
|
+
Provides-Extra: watch
|
|
21
|
+
Requires-Dist: watchfiles>=0.24; extra == 'watch'
|
|
22
|
+
Description-Content-Type: text/markdown
|
|
23
|
+
|
|
24
|
+
# torchtyc
|
|
25
|
+
|
|
26
|
+
Static array shape checking for PyTorch, powered by meta tensors.
|
|
27
|
+
|
|
28
|
+
torchtyc reads [jaxtyping](https://docs.kidger.site/jaxtyping/) annotations and
|
|
29
|
+
verifies the shapes your code actually produces, before you run it on a GPU. It
|
|
30
|
+
is the PyTorch counterpart to [jaxtyc](https://github.com/BeeGass/jaxtyc), which
|
|
31
|
+
does the same thing for JAX with `jax.eval_shape`.
|
|
32
|
+
|
|
33
|
+
```python
|
|
34
|
+
# model.py
|
|
35
|
+
import torch
|
|
36
|
+
from einops import einsum
|
|
37
|
+
from jaxtyping import Float
|
|
38
|
+
from torch import Tensor, nn
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
class Linear(nn.Module):
|
|
42
|
+
def __init__(self, in_features: int, out_features: int) -> None:
|
|
43
|
+
super().__init__()
|
|
44
|
+
self.W = nn.Parameter(torch.empty((out_features, in_features)))
|
|
45
|
+
|
|
46
|
+
def forward(self, x: Float[Tensor, "... in_features"]) -> Float[Tensor, "... in_features"]:
|
|
47
|
+
return einsum(x, self.W, "... in_features, out_features in_features -> ... out_features")
|
|
48
|
+
```
|
|
49
|
+
|
|
50
|
+
```
|
|
51
|
+
$ torchtyc check model.py
|
|
52
|
+
model.py:12:63: error[shape-mismatch]
|
|
53
|
+
in the return of `Linear.forward`: annotated `in_features`, but the traced dimension is `out_features`
|
|
54
|
+
Expected: (..., in_features)
|
|
55
|
+
Got: (..., out_features)
|
|
56
|
+
12 | def forward(self, x: Float[Tensor, "... in_features"]) -> Float[Tensor, "... in_features"]:
|
|
57
|
+
try: Float[Tensor, "... out_features"]
|
|
58
|
+
hint: this dimension is `out_features`, so the annotation likely names the wrong axis
|
|
59
|
+
|
|
60
|
+
Found 1 error(s) in 1 function(s) across 1 file(s)
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
Pyright and mypy cannot do this. To them `Float[Tensor, "... in_features"]` is
|
|
64
|
+
just `Tensor`, and the dim string is an opaque literal.
|
|
65
|
+
|
|
66
|
+
## How it works
|
|
67
|
+
|
|
68
|
+
torchtyc constructs each annotated function's arguments on
|
|
69
|
+
`torch.device("meta")`, calls the function, and compares the shape that comes
|
|
70
|
+
back against the annotation. Meta tensors carry shape, dtype, and stride but own
|
|
71
|
+
no storage, so a whole model runs for the cost of a dictionary lookup per
|
|
72
|
+
operator, with no allocation and no arithmetic.
|
|
73
|
+
|
|
74
|
+
Every dimension name is bound to a distinct prime number starting at 101. This
|
|
75
|
+
buys two things:
|
|
76
|
+
|
|
77
|
+
**Mistakes cannot hide.** If `d_in` and `d_out` were both bound to 64, a
|
|
78
|
+
transposed weight matrix would sail through. Distinct primes make the two
|
|
79
|
+
impossible to confuse.
|
|
80
|
+
|
|
81
|
+
**Products stay readable.** A traced dimension that is the product of two
|
|
82
|
+
bound primes factors back into the names that produced it, which tells you the
|
|
83
|
+
function flattened two axes together. The primes themselves never reach the
|
|
84
|
+
message:
|
|
85
|
+
|
|
86
|
+
```python
|
|
87
|
+
# flat.py
|
|
88
|
+
def flatten(x: Float[Tensor, "batch seq d_model"]) -> Float[Tensor, "batch seq d_model"]:
|
|
89
|
+
return x.reshape(x.shape[0], -1)
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
```
|
|
93
|
+
flat.py:6:55: error[rank-mismatch]
|
|
94
|
+
in the return of `flatten`: expected 3 dimensions, traced 2
|
|
95
|
+
Expected: (batch, seq, d_model)
|
|
96
|
+
Got: (batch, d_model*seq)
|
|
97
|
+
6 | def flatten(x: Float[Tensor, "batch seq d_model"]) -> Float[Tensor, "batch seq d_model"]:
|
|
98
|
+
hint: d_model*seq looks like two annotated axes flattened into one
|
|
99
|
+
```
|
|
100
|
+
|
|
101
|
+
An axis you did not name renders as `...`, a single `_` axis renders as `_`,
|
|
102
|
+
and a raw number appears only for a size the annotation never bound.
|
|
103
|
+
|
|
104
|
+
Because it runs your function rather than reasoning about it symbolically,
|
|
105
|
+
torchtyc also catches anything that raises on the way: a bad `einsum`, a
|
|
106
|
+
`matmul` between incompatible operands, an `nn.Module` that cannot be built. The
|
|
107
|
+
diagnostic anchors to the deepest frame inside your own file, not to a line in
|
|
108
|
+
torch.
|
|
109
|
+
|
|
110
|
+
## Install
|
|
111
|
+
|
|
112
|
+
```bash
|
|
113
|
+
uv add --dev torchtyc # or: pip install torchtyc
|
|
114
|
+
```
|
|
115
|
+
|
|
116
|
+
Extras: `torchtyc[lsp]` for the language server, `[watch]` for watch mode,
|
|
117
|
+
`[einops]` for einops-aware hints, `[all]` for everything.
|
|
118
|
+
|
|
119
|
+
torchtyc must run under the same interpreter as your project, since it imports
|
|
120
|
+
your code. By default it finds `.venv/bin/python` next to your `pyproject.toml`.
|
|
121
|
+
Override with `--python` or `[tool.torchtyc] python = "..."`.
|
|
122
|
+
|
|
123
|
+
## Commands
|
|
124
|
+
|
|
125
|
+
```
|
|
126
|
+
torchtyc check <paths>... Shape-check files or directories
|
|
127
|
+
torchtyc trace <file.py::func> Show the shapes flowing through one function
|
|
128
|
+
torchtyc watch <paths>... Re-check on change
|
|
129
|
+
torchtyc lsp Language server on stdio
|
|
130
|
+
torchtyc lsp --tcp PORT Language server on a TCP port
|
|
131
|
+
torchtyc rules List the diagnostic rules
|
|
132
|
+
torchtyc version Print the version
|
|
133
|
+
```
|
|
134
|
+
|
|
135
|
+
`check` takes `--format full` (default), `concise`, `json`, or `github`.
|
|
136
|
+
|
|
137
|
+
```
|
|
138
|
+
$ torchtyc trace model.py::Linear.forward
|
|
139
|
+
Linear.forward
|
|
140
|
+
x : (..., in_features)
|
|
141
|
+
return -> float32[(..., out_features)]
|
|
142
|
+
|
|
143
|
+
dimension names are bound to distinct primes starting at 101
|
|
144
|
+
```
|
|
145
|
+
|
|
146
|
+
## Constructing modules
|
|
147
|
+
|
|
148
|
+
To check a method, torchtyc needs an instance, and to build an instance it needs
|
|
149
|
+
constructor arguments. It matches integer parameters of `__init__` against the
|
|
150
|
+
dimension names in the method's annotations:
|
|
151
|
+
|
|
152
|
+
```python
|
|
153
|
+
def __init__(self, in_features: int, out_features: int) -> None: ...
|
|
154
|
+
def forward(self, x: Float[Tensor, "... in_features"]) -> ...
|
|
155
|
+
```
|
|
156
|
+
|
|
157
|
+
`in_features` is a dimension name, so it receives that dimension's prime.
|
|
158
|
+
Parameters with defaults are left alone. A parameter that is neither a
|
|
159
|
+
dimension, a known type, nor defaulted produces an `unresolved-arg` warning and
|
|
160
|
+
the function is skipped rather than guessed at.
|
|
161
|
+
|
|
162
|
+
Modules are built inside `torch.device("meta")`, and the initialisers in
|
|
163
|
+
`torch.nn.init` are neutralised for the duration, since initial values cannot
|
|
164
|
+
affect a shape.
|
|
165
|
+
|
|
166
|
+
## Annotated attributes
|
|
167
|
+
|
|
168
|
+
Python does not check variable annotations at runtime, and neither does
|
|
169
|
+
jaxtyping, which only reads function signatures. Since torchtyc has a
|
|
170
|
+
constructed instance in hand anyway, it checks them too:
|
|
171
|
+
|
|
172
|
+
```python
|
|
173
|
+
# attr.py
|
|
174
|
+
class Linear(nn.Module):
|
|
175
|
+
def __init__(self, d_in: int, d_out: int) -> None:
|
|
176
|
+
super().__init__()
|
|
177
|
+
self.W: Float[nn.Parameter, "d_out d_in"] = nn.Parameter(torch.empty((d_in, d_out)))
|
|
178
|
+
```
|
|
179
|
+
|
|
180
|
+
```
|
|
181
|
+
attr.py:9:17: error[attribute-mismatch]
|
|
182
|
+
`self.W`: annotated `d_out`, but the traced dimension is `d_in`
|
|
183
|
+
Expected: (d_out, d_in)
|
|
184
|
+
Got: (d_in, d_out)
|
|
185
|
+
9 | self.W: Float[nn.Parameter, "d_out d_in"] = nn.Parameter(torch.empty((d_in, d_out)))
|
|
186
|
+
try: Float[nn.Parameter, "d_in d_out"]
|
|
187
|
+
hint: this dimension is `d_in`, so the annotation likely names the wrong axis
|
|
188
|
+
```
|
|
189
|
+
|
|
190
|
+
## Suppressing
|
|
191
|
+
|
|
192
|
+
```python
|
|
193
|
+
y = x.reshape(-1) # torchtyc: ignore
|
|
194
|
+
y = x.reshape(-1) # torchtyc: ignore[rank-mismatch]
|
|
195
|
+
```
|
|
196
|
+
|
|
197
|
+
A scoped ignore that never matches is itself reported, so suppressions do not
|
|
198
|
+
rot.
|
|
199
|
+
|
|
200
|
+
## Configuration
|
|
201
|
+
|
|
202
|
+
```toml
|
|
203
|
+
[tool.torchtyc]
|
|
204
|
+
python = ".venv/bin/python" # interpreter that imports your code
|
|
205
|
+
severity = "warning" # drop anything below this level
|
|
206
|
+
ignore = ["unused-dim"]
|
|
207
|
+
exclude = [".venv", "build", "experiments"]
|
|
208
|
+
variadic-rank = 2 # how many axes `...` stands for
|
|
209
|
+
einops = true
|
|
210
|
+
timeout = 60.0
|
|
211
|
+
extra-paths = ["stubs"] # prepended to the worker's PYTHONPATH
|
|
212
|
+
```
|
|
213
|
+
|
|
214
|
+
## Editors
|
|
215
|
+
|
|
216
|
+
Any LSP client works. Neovim, without a plugin:
|
|
217
|
+
|
|
218
|
+
```lua
|
|
219
|
+
vim.lsp.config.torchtyc = {
|
|
220
|
+
cmd = { "torchtyc", "lsp" },
|
|
221
|
+
filetypes = { "python" },
|
|
222
|
+
root_markers = { "pyproject.toml", ".git" },
|
|
223
|
+
}
|
|
224
|
+
vim.lsp.enable("torchtyc")
|
|
225
|
+
```
|
|
226
|
+
|
|
227
|
+
The server publishes lint diagnostics immediately on every change, and traces
|
|
228
|
+
after the buffer has been quiet for 0.7s, on open, and on save. It never imports
|
|
229
|
+
your code on a keystroke.
|
|
230
|
+
|
|
231
|
+
Hover over a function to see the traced shapes. Inlay hints show the traced
|
|
232
|
+
return next to each signature, and a code lens above it reports the traced
|
|
233
|
+
return or the error count. Code actions offer to silence a rule or to adopt
|
|
234
|
+
the shape that was actually traced.
|
|
235
|
+
|
|
236
|
+
## CI
|
|
237
|
+
|
|
238
|
+
```yaml
|
|
239
|
+
- run: uv run torchtyc check src/ --format github
|
|
240
|
+
```
|
|
241
|
+
|
|
242
|
+
`--format github` emits workflow commands, so each finding becomes an inline
|
|
243
|
+
annotation on the pull request diff. `check` exits 1 when there are errors, and
|
|
244
|
+
2 when torchtyc itself could not do the job: the worker failed, or the paths
|
|
245
|
+
matched no python file.
|
|
246
|
+
|
|
247
|
+
## Rules
|
|
248
|
+
|
|
249
|
+
| Rule | Level | Meaning |
|
|
250
|
+
| --- | --- | --- |
|
|
251
|
+
| `shape-mismatch` | error | a traced shape disagrees with its annotation |
|
|
252
|
+
| `rank-mismatch` | error | a traced value has a different number of dimensions |
|
|
253
|
+
| `dtype-mismatch` | error | a traced dtype is outside the annotated dtype set |
|
|
254
|
+
| `dim-inconsistent` | error | one dimension name is bound to two different sizes |
|
|
255
|
+
| `attribute-mismatch` | error | an annotated attribute on self holds a different shape |
|
|
256
|
+
| `not-a-tensor` | error | an annotated tensor position received a non-tensor |
|
|
257
|
+
| `tuple-arity` | error | a tuple return has a different length than annotated |
|
|
258
|
+
| `einops-pattern` | error | an einops pattern disagrees with the tensors given to it |
|
|
259
|
+
| `trace-error` | error | the function raised while being traced |
|
|
260
|
+
| `import-error` | error | the module could not be imported |
|
|
261
|
+
| `device-mismatch` | warning | a traced value left the meta device |
|
|
262
|
+
| `einops-unknown-axis` | warning | an einops axis matches no input axis or keyword |
|
|
263
|
+
| `uninstantiable` | warning | a module's `__init__` could not be called automatically |
|
|
264
|
+
| `unresolved-arg` | warning | a parameter has no annotation and no default |
|
|
265
|
+
| `unsupported-annotation` | warning | an annotation could not be parsed |
|
|
266
|
+
| `local-definition` | info | a target inside a function body cannot be reached after import |
|
|
267
|
+
| `anonymous-return` | info | arguments are annotated but the return is not |
|
|
268
|
+
| `missing-annotation` | info | a public function has no jaxtyping annotation |
|
|
269
|
+
| `unused-dim` | info | a dimension name is used once, so it constrains nothing |
|
|
270
|
+
| `suppression-unused` | info | an ignore comment matched no diagnostic |
|
|
271
|
+
|
|
272
|
+
## Limits
|
|
273
|
+
|
|
274
|
+
torchtyc imports your module, so module-level side effects run. Keep training
|
|
275
|
+
loops behind `if __name__ == "__main__":`.
|
|
276
|
+
|
|
277
|
+
Classes and functions nested inside other classes, or written under a
|
|
278
|
+
module-level `if` or `try`, are checked like any other. One inside a *function
|
|
279
|
+
body* is not: after import there is no name to reach it by, so it reports
|
|
280
|
+
`local-definition` rather than passing silently. A `if TYPE_CHECKING:` block is
|
|
281
|
+
skipped without a diagnostic, because nothing in it exists at runtime.
|
|
282
|
+
|
|
283
|
+
It runs one concrete trace, not a proof. A function whose control flow depends
|
|
284
|
+
on tensor *values* rather than shapes takes whichever branch the primes send it
|
|
285
|
+
down. `...` stands for a fixed number of axes, two by default, so code that
|
|
286
|
+
behaves differently at other ranks needs `variadic-rank` or a second annotated
|
|
287
|
+
wrapper. Every bare `...` in one signature also stands for the *same* axes,
|
|
288
|
+
which is narrower than jaxtyping: a loss taking two `"... d"` arguments traces
|
|
289
|
+
with one batch shape, because that is what the annotation almost always means.
|
|
290
|
+
|
|
291
|
+
A dimension is a prime, and a prime does not divide. Code that splits an axis
|
|
292
|
+
- `head_dim = d_model // n_heads`, then `view(b, s, n_heads, head_dim)` - gets a
|
|
293
|
+
quotient that does not multiply back, so correct multi-head attention is
|
|
294
|
+
reported as a `trace-error`. A default on the divided parameter does not help,
|
|
295
|
+
because a name the annotations use is bound to its prime whatever it defaults
|
|
296
|
+
to. Silence the line with `# torchtyc: ignore[trace-error]` until dimensions
|
|
297
|
+
carry a divisible factor.
|
|
298
|
+
|
|
299
|
+
Runtime checking with `jaxtyping` and `beartype` remains worth having. torchtyc
|
|
300
|
+
tells you the shapes are consistent for the sizes it chose; beartype tells you
|
|
301
|
+
they were right for the batch you actually ran.
|
|
302
|
+
|
|
303
|
+
## Licence
|
|
304
|
+
|
|
305
|
+
MIT.
|
torchtyc-0.1.0/README.md
ADDED
|
@@ -0,0 +1,282 @@
|
|
|
1
|
+
# torchtyc
|
|
2
|
+
|
|
3
|
+
Static array shape checking for PyTorch, powered by meta tensors.
|
|
4
|
+
|
|
5
|
+
torchtyc reads [jaxtyping](https://docs.kidger.site/jaxtyping/) annotations and
|
|
6
|
+
verifies the shapes your code actually produces, before you run it on a GPU. It
|
|
7
|
+
is the PyTorch counterpart to [jaxtyc](https://github.com/BeeGass/jaxtyc), which
|
|
8
|
+
does the same thing for JAX with `jax.eval_shape`.
|
|
9
|
+
|
|
10
|
+
```python
|
|
11
|
+
# model.py
|
|
12
|
+
import torch
|
|
13
|
+
from einops import einsum
|
|
14
|
+
from jaxtyping import Float
|
|
15
|
+
from torch import Tensor, nn
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class Linear(nn.Module):
|
|
19
|
+
def __init__(self, in_features: int, out_features: int) -> None:
|
|
20
|
+
super().__init__()
|
|
21
|
+
self.W = nn.Parameter(torch.empty((out_features, in_features)))
|
|
22
|
+
|
|
23
|
+
def forward(self, x: Float[Tensor, "... in_features"]) -> Float[Tensor, "... in_features"]:
|
|
24
|
+
return einsum(x, self.W, "... in_features, out_features in_features -> ... out_features")
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
```
|
|
28
|
+
$ torchtyc check model.py
|
|
29
|
+
model.py:12:63: error[shape-mismatch]
|
|
30
|
+
in the return of `Linear.forward`: annotated `in_features`, but the traced dimension is `out_features`
|
|
31
|
+
Expected: (..., in_features)
|
|
32
|
+
Got: (..., out_features)
|
|
33
|
+
12 | def forward(self, x: Float[Tensor, "... in_features"]) -> Float[Tensor, "... in_features"]:
|
|
34
|
+
try: Float[Tensor, "... out_features"]
|
|
35
|
+
hint: this dimension is `out_features`, so the annotation likely names the wrong axis
|
|
36
|
+
|
|
37
|
+
Found 1 error(s) in 1 function(s) across 1 file(s)
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
Pyright and mypy cannot do this. To them `Float[Tensor, "... in_features"]` is
|
|
41
|
+
just `Tensor`, and the dim string is an opaque literal.
|
|
42
|
+
|
|
43
|
+
## How it works
|
|
44
|
+
|
|
45
|
+
torchtyc constructs each annotated function's arguments on
|
|
46
|
+
`torch.device("meta")`, calls the function, and compares the shape that comes
|
|
47
|
+
back against the annotation. Meta tensors carry shape, dtype, and stride but own
|
|
48
|
+
no storage, so a whole model runs for the cost of a dictionary lookup per
|
|
49
|
+
operator, with no allocation and no arithmetic.
|
|
50
|
+
|
|
51
|
+
Every dimension name is bound to a distinct prime number starting at 101. This
|
|
52
|
+
buys two things:
|
|
53
|
+
|
|
54
|
+
**Mistakes cannot hide.** If `d_in` and `d_out` were both bound to 64, a
|
|
55
|
+
transposed weight matrix would sail through. Distinct primes make the two
|
|
56
|
+
impossible to confuse.
|
|
57
|
+
|
|
58
|
+
**Products stay readable.** A traced dimension that is the product of two
|
|
59
|
+
bound primes factors back into the names that produced it, which tells you the
|
|
60
|
+
function flattened two axes together. The primes themselves never reach the
|
|
61
|
+
message:
|
|
62
|
+
|
|
63
|
+
```python
|
|
64
|
+
# flat.py
|
|
65
|
+
def flatten(x: Float[Tensor, "batch seq d_model"]) -> Float[Tensor, "batch seq d_model"]:
|
|
66
|
+
return x.reshape(x.shape[0], -1)
|
|
67
|
+
```
|
|
68
|
+
|
|
69
|
+
```
|
|
70
|
+
flat.py:6:55: error[rank-mismatch]
|
|
71
|
+
in the return of `flatten`: expected 3 dimensions, traced 2
|
|
72
|
+
Expected: (batch, seq, d_model)
|
|
73
|
+
Got: (batch, d_model*seq)
|
|
74
|
+
6 | def flatten(x: Float[Tensor, "batch seq d_model"]) -> Float[Tensor, "batch seq d_model"]:
|
|
75
|
+
hint: d_model*seq looks like two annotated axes flattened into one
|
|
76
|
+
```
|
|
77
|
+
|
|
78
|
+
An axis you did not name renders as `...`, a single `_` axis renders as `_`,
|
|
79
|
+
and a raw number appears only for a size the annotation never bound.
|
|
80
|
+
|
|
81
|
+
Because it runs your function rather than reasoning about it symbolically,
|
|
82
|
+
torchtyc also catches anything that raises on the way: a bad `einsum`, a
|
|
83
|
+
`matmul` between incompatible operands, an `nn.Module` that cannot be built. The
|
|
84
|
+
diagnostic anchors to the deepest frame inside your own file, not to a line in
|
|
85
|
+
torch.
|
|
86
|
+
|
|
87
|
+
## Install
|
|
88
|
+
|
|
89
|
+
```bash
|
|
90
|
+
uv add --dev torchtyc # or: pip install torchtyc
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
Extras: `torchtyc[lsp]` for the language server, `[watch]` for watch mode,
|
|
94
|
+
`[einops]` for einops-aware hints, `[all]` for everything.
|
|
95
|
+
|
|
96
|
+
torchtyc must run under the same interpreter as your project, since it imports
|
|
97
|
+
your code. By default it finds `.venv/bin/python` next to your `pyproject.toml`.
|
|
98
|
+
Override with `--python` or `[tool.torchtyc] python = "..."`.
|
|
99
|
+
|
|
100
|
+
## Commands
|
|
101
|
+
|
|
102
|
+
```
|
|
103
|
+
torchtyc check <paths>... Shape-check files or directories
|
|
104
|
+
torchtyc trace <file.py::func> Show the shapes flowing through one function
|
|
105
|
+
torchtyc watch <paths>... Re-check on change
|
|
106
|
+
torchtyc lsp Language server on stdio
|
|
107
|
+
torchtyc lsp --tcp PORT Language server on a TCP port
|
|
108
|
+
torchtyc rules List the diagnostic rules
|
|
109
|
+
torchtyc version Print the version
|
|
110
|
+
```
|
|
111
|
+
|
|
112
|
+
`check` takes `--format full` (default), `concise`, `json`, or `github`.
|
|
113
|
+
|
|
114
|
+
```
|
|
115
|
+
$ torchtyc trace model.py::Linear.forward
|
|
116
|
+
Linear.forward
|
|
117
|
+
x : (..., in_features)
|
|
118
|
+
return -> float32[(..., out_features)]
|
|
119
|
+
|
|
120
|
+
dimension names are bound to distinct primes starting at 101
|
|
121
|
+
```
|
|
122
|
+
|
|
123
|
+
## Constructing modules
|
|
124
|
+
|
|
125
|
+
To check a method, torchtyc needs an instance, and to build an instance it needs
|
|
126
|
+
constructor arguments. It matches integer parameters of `__init__` against the
|
|
127
|
+
dimension names in the method's annotations:
|
|
128
|
+
|
|
129
|
+
```python
|
|
130
|
+
def __init__(self, in_features: int, out_features: int) -> None: ...
|
|
131
|
+
def forward(self, x: Float[Tensor, "... in_features"]) -> ...
|
|
132
|
+
```
|
|
133
|
+
|
|
134
|
+
`in_features` is a dimension name, so it receives that dimension's prime.
|
|
135
|
+
Parameters with defaults are left alone. A parameter that is neither a
|
|
136
|
+
dimension, a known type, nor defaulted produces an `unresolved-arg` warning and
|
|
137
|
+
the function is skipped rather than guessed at.
|
|
138
|
+
|
|
139
|
+
Modules are built inside `torch.device("meta")`, and the initialisers in
|
|
140
|
+
`torch.nn.init` are neutralised for the duration, since initial values cannot
|
|
141
|
+
affect a shape.
|
|
142
|
+
|
|
143
|
+
## Annotated attributes
|
|
144
|
+
|
|
145
|
+
Python does not check variable annotations at runtime, and neither does
|
|
146
|
+
jaxtyping, which only reads function signatures. Since torchtyc has a
|
|
147
|
+
constructed instance in hand anyway, it checks them too:
|
|
148
|
+
|
|
149
|
+
```python
|
|
150
|
+
# attr.py
|
|
151
|
+
class Linear(nn.Module):
|
|
152
|
+
def __init__(self, d_in: int, d_out: int) -> None:
|
|
153
|
+
super().__init__()
|
|
154
|
+
self.W: Float[nn.Parameter, "d_out d_in"] = nn.Parameter(torch.empty((d_in, d_out)))
|
|
155
|
+
```
|
|
156
|
+
|
|
157
|
+
```
|
|
158
|
+
attr.py:9:17: error[attribute-mismatch]
|
|
159
|
+
`self.W`: annotated `d_out`, but the traced dimension is `d_in`
|
|
160
|
+
Expected: (d_out, d_in)
|
|
161
|
+
Got: (d_in, d_out)
|
|
162
|
+
9 | self.W: Float[nn.Parameter, "d_out d_in"] = nn.Parameter(torch.empty((d_in, d_out)))
|
|
163
|
+
try: Float[nn.Parameter, "d_in d_out"]
|
|
164
|
+
hint: this dimension is `d_in`, so the annotation likely names the wrong axis
|
|
165
|
+
```
|
|
166
|
+
|
|
167
|
+
## Suppressing
|
|
168
|
+
|
|
169
|
+
```python
|
|
170
|
+
y = x.reshape(-1) # torchtyc: ignore
|
|
171
|
+
y = x.reshape(-1) # torchtyc: ignore[rank-mismatch]
|
|
172
|
+
```
|
|
173
|
+
|
|
174
|
+
A scoped ignore that never matches is itself reported, so suppressions do not
|
|
175
|
+
rot.
|
|
176
|
+
|
|
177
|
+
## Configuration
|
|
178
|
+
|
|
179
|
+
```toml
|
|
180
|
+
[tool.torchtyc]
|
|
181
|
+
python = ".venv/bin/python" # interpreter that imports your code
|
|
182
|
+
severity = "warning" # drop anything below this level
|
|
183
|
+
ignore = ["unused-dim"]
|
|
184
|
+
exclude = [".venv", "build", "experiments"]
|
|
185
|
+
variadic-rank = 2 # how many axes `...` stands for
|
|
186
|
+
einops = true
|
|
187
|
+
timeout = 60.0
|
|
188
|
+
extra-paths = ["stubs"] # prepended to the worker's PYTHONPATH
|
|
189
|
+
```
|
|
190
|
+
|
|
191
|
+
## Editors
|
|
192
|
+
|
|
193
|
+
Any LSP client works. Neovim, without a plugin:
|
|
194
|
+
|
|
195
|
+
```lua
|
|
196
|
+
vim.lsp.config.torchtyc = {
|
|
197
|
+
cmd = { "torchtyc", "lsp" },
|
|
198
|
+
filetypes = { "python" },
|
|
199
|
+
root_markers = { "pyproject.toml", ".git" },
|
|
200
|
+
}
|
|
201
|
+
vim.lsp.enable("torchtyc")
|
|
202
|
+
```
|
|
203
|
+
|
|
204
|
+
The server publishes lint diagnostics immediately on every change, and traces
|
|
205
|
+
after the buffer has been quiet for 0.7s, on open, and on save. It never imports
|
|
206
|
+
your code on a keystroke.
|
|
207
|
+
|
|
208
|
+
Hover over a function to see the traced shapes. Inlay hints show the traced
|
|
209
|
+
return next to each signature, and a code lens above it reports the traced
|
|
210
|
+
return or the error count. Code actions offer to silence a rule or to adopt
|
|
211
|
+
the shape that was actually traced.
|
|
212
|
+
|
|
213
|
+
## CI
|
|
214
|
+
|
|
215
|
+
```yaml
|
|
216
|
+
- run: uv run torchtyc check src/ --format github
|
|
217
|
+
```
|
|
218
|
+
|
|
219
|
+
`--format github` emits workflow commands, so each finding becomes an inline
|
|
220
|
+
annotation on the pull request diff. `check` exits 1 when there are errors, and
|
|
221
|
+
2 when torchtyc itself could not do the job: the worker failed, or the paths
|
|
222
|
+
matched no python file.
|
|
223
|
+
|
|
224
|
+
## Rules
|
|
225
|
+
|
|
226
|
+
| Rule | Level | Meaning |
|
|
227
|
+
| --- | --- | --- |
|
|
228
|
+
| `shape-mismatch` | error | a traced shape disagrees with its annotation |
|
|
229
|
+
| `rank-mismatch` | error | a traced value has a different number of dimensions |
|
|
230
|
+
| `dtype-mismatch` | error | a traced dtype is outside the annotated dtype set |
|
|
231
|
+
| `dim-inconsistent` | error | one dimension name is bound to two different sizes |
|
|
232
|
+
| `attribute-mismatch` | error | an annotated attribute on self holds a different shape |
|
|
233
|
+
| `not-a-tensor` | error | an annotated tensor position received a non-tensor |
|
|
234
|
+
| `tuple-arity` | error | a tuple return has a different length than annotated |
|
|
235
|
+
| `einops-pattern` | error | an einops pattern disagrees with the tensors given to it |
|
|
236
|
+
| `trace-error` | error | the function raised while being traced |
|
|
237
|
+
| `import-error` | error | the module could not be imported |
|
|
238
|
+
| `device-mismatch` | warning | a traced value left the meta device |
|
|
239
|
+
| `einops-unknown-axis` | warning | an einops axis matches no input axis or keyword |
|
|
240
|
+
| `uninstantiable` | warning | a module's `__init__` could not be called automatically |
|
|
241
|
+
| `unresolved-arg` | warning | a parameter has no annotation and no default |
|
|
242
|
+
| `unsupported-annotation` | warning | an annotation could not be parsed |
|
|
243
|
+
| `local-definition` | info | a target inside a function body cannot be reached after import |
|
|
244
|
+
| `anonymous-return` | info | arguments are annotated but the return is not |
|
|
245
|
+
| `missing-annotation` | info | a public function has no jaxtyping annotation |
|
|
246
|
+
| `unused-dim` | info | a dimension name is used once, so it constrains nothing |
|
|
247
|
+
| `suppression-unused` | info | an ignore comment matched no diagnostic |
|
|
248
|
+
|
|
249
|
+
## Limits
|
|
250
|
+
|
|
251
|
+
torchtyc imports your module, so module-level side effects run. Keep training
|
|
252
|
+
loops behind `if __name__ == "__main__":`.
|
|
253
|
+
|
|
254
|
+
Classes and functions nested inside other classes, or written under a
|
|
255
|
+
module-level `if` or `try`, are checked like any other. One inside a *function
|
|
256
|
+
body* is not: after import there is no name to reach it by, so it reports
|
|
257
|
+
`local-definition` rather than passing silently. A `if TYPE_CHECKING:` block is
|
|
258
|
+
skipped without a diagnostic, because nothing in it exists at runtime.
|
|
259
|
+
|
|
260
|
+
It runs one concrete trace, not a proof. A function whose control flow depends
|
|
261
|
+
on tensor *values* rather than shapes takes whichever branch the primes send it
|
|
262
|
+
down. `...` stands for a fixed number of axes, two by default, so code that
|
|
263
|
+
behaves differently at other ranks needs `variadic-rank` or a second annotated
|
|
264
|
+
wrapper. Every bare `...` in one signature also stands for the *same* axes,
|
|
265
|
+
which is narrower than jaxtyping: a loss taking two `"... d"` arguments traces
|
|
266
|
+
with one batch shape, because that is what the annotation almost always means.
|
|
267
|
+
|
|
268
|
+
A dimension is a prime, and a prime does not divide. Code that splits an axis
|
|
269
|
+
- `head_dim = d_model // n_heads`, then `view(b, s, n_heads, head_dim)` - gets a
|
|
270
|
+
quotient that does not multiply back, so correct multi-head attention is
|
|
271
|
+
reported as a `trace-error`. A default on the divided parameter does not help,
|
|
272
|
+
because a name the annotations use is bound to its prime whatever it defaults
|
|
273
|
+
to. Silence the line with `# torchtyc: ignore[trace-error]` until dimensions
|
|
274
|
+
carry a divisible factor.
|
|
275
|
+
|
|
276
|
+
Runtime checking with `jaxtyping` and `beartype` remains worth having. torchtyc
|
|
277
|
+
tells you the shapes are consistent for the sizes it chose; beartype tells you
|
|
278
|
+
they were right for the batch you actually ran.
|
|
279
|
+
|
|
280
|
+
## Licence
|
|
281
|
+
|
|
282
|
+
MIT.
|