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.
@@ -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
@@ -0,0 +1,11 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ .venv/
4
+ dist/
5
+ build/
6
+ *.egg-info/
7
+ .pytest_cache/
8
+ .ruff_cache/
9
+ result
10
+ uv.lock
11
+ .env
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.
@@ -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.
@@ -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.