torchtyc

Shape checking for PyTorch, from your jaxtyping annotations.

torchtyc reads jaxtyping annotations and verifies the shapes your code actually produces, before you run it on a GPU. It is the PyTorch counterpart to jaxtyc, which does the same thing for JAX with jax.eval_shape.

# model.py
import torch
from einops import einsum
from jaxtyping import Float
from torch import Tensor, nn


class Linear(nn.Module):
    def __init__(self, in_features: int, out_features: int) -> None:
        super().__init__()
        self.W = nn.Parameter(torch.empty((out_features, in_features)))

    def forward(self, x: Float[Tensor, "... in_features"]) -> Float[Tensor, "... in_features"]:
        return einsum(x, self.W, "... in_features, out_features in_features -> ... out_features")

$ torchtyc check model.py
model.py:12:63: error[shape-mismatch]
  in the return of `Linear.forward`: annotated `in_features`, but the traced dimension is `out_features`
    Expected: (..., in_features)
    Got:      (..., out_features)
    12 | def forward(self, x: Float[Tensor, "... in_features"]) -> Float[Tensor, "... in_features"]:
  try:  Float[Tensor, "... out_features"]
  hint: this dimension is `out_features`, so the annotation likely names the wrong axis

Found 1 error(s) in 1 function(s) across 1 file(s)

Pyright and mypy cannot do this. To them Float[Tensor, "... in_features"] is just Tensor, and the dim string is an opaque literal.

How it works

torchtyc constructs each annotated function's arguments on torch.device("meta"), calls the function, and compares the shape that comes back against the annotation. Meta tensors carry shape, dtype, and stride but own no storage, so a whole model runs for the cost of a dictionary lookup per operator, with no allocation and no arithmetic.

Every dimension name is bound to a distinct prime number starting at 101. This buys two things:

Mistakes cannot hide. If d_in and d_out were both bound to 64, a transposed weight matrix would sail through. Distinct primes make the two impossible to confuse.

Products stay readable. A traced dimension that is the product of two bound primes factors back into the names that produced it, which tells you the function flattened two axes together. The primes themselves never reach the message:

# flat.py
def flatten(x: Float[Tensor, "batch seq d_model"]) -> Float[Tensor, "batch seq d_model"]:
    return x.reshape(x.shape[0], -1)

flat.py:6:55: error[rank-mismatch]
  in the return of `flatten`: expected 3 dimensions, traced 2
    Expected: (batch, seq, d_model)
    Got:      (batch, d_model*seq)
    6 | def flatten(x: Float[Tensor, "batch seq d_model"]) -> Float[Tensor, "batch seq d_model"]:
  hint: d_model*seq looks like two annotated axes flattened into one

An axis you did not name renders as ..., and a single _ axis renders as _. A width your own __init__ worked out renders as what it follows, such as <from d_model>, because the number it lands on comes from the width torchtyc traced with and is not a width your model has. A plain number appears only for a size your code writes down, which is one you can go and find.

Because it runs your function rather than reasoning about it symbolically, torchtyc also catches anything that raises on the way: a bad einsum, a matmul between incompatible operands, an nn.Module that cannot be built.

layers.py:129:45: error[trace-error]
  RuntimeError: einsum(): subscript a has size <from d_model> for operand 1 which does not broadcast with previously seen size d_model
    129 | w1_out: Float[Tensor, "... d_ff"] = self.w1(x)
  hint: raised further down, in `Linear.forward` at line 40, where x is (..., d_model), self.weight is (d_model, <from d_model>)
  note: <from d_model> is a width your __init__ computed from d_model, so it is not d_model

The diagnostic anchors to the line you wrote, not to a line in torch and not to the shared layer further down that every caller runs. The hint carries the shapes of the tensors on the line that raised, so both widths that disagree are on show: here self.w1 was built with its two widths the wrong way round.

Install

Into the project, so the version is pinned alongside everything else:

uv add --dev "torchtyc[all]"      # or: pip install "torchtyc[all]"

Or once, for every project on the machine:

uv tool install "torchtyc[all]"

Extras: torchtyc[lsp] for the language server, [watch] for watch mode, [einops] for einops-aware hints, [all] for everything.

torchtyc itself runs under any interpreter. Your code is imported by a separate worker process, which torchtyc starts under your project's interpreter with its own source added to that interpreter's path. A standalone torchtyc on 3.14 therefore checks a project on 3.13 without being installed into that project's venv.

The worker needs torch importable, plus whatever your modules import at import time. torchtyc looks for .venv/bin/python, then venv/bin/python, beside the nearest pyproject.toml above the paths you gave it, and falls back to the interpreter it is running under. Override with --python or [tool.torchtyc] python = "...".

Confirm the pieces line up before anything else:

torchtyc version
torchtyc check src/one_model.py

A worker failed: line means the interpreter is wrong or your module does not import. Anything else is a finding about your code.

Checking every file

torchtyc check .

A directory argument is walked recursively for *.py. .venv, build, dist, __pycache__ and .git are skipped. Setting exclude in pyproject.toml replaces that list, so repeat the defaults you still want:

[tool.torchtyc]
exclude = [".venv", "build", "dist", "__pycache__", ".git", "experiments"]

Each name matches a whole path component below the directory you named, so experiments drops src/experiments/train.py. Directories above the one you asked for are not considered, which is why a checkout that happens to live under a path called build still gets checked when you point at it directly.

The language server only ever checks the buffer you have open. Project-wide coverage comes from the CLI. While you work:

torchtyc watch .

Before a commit, as a .pre-commit-config.yaml hook:

- repo: local
  hooks:
    - id: torchtyc
      name: torchtyc
      entry: uv run torchtyc check
      language: system
      types: [python]

On a pull request, see CI below.

Commands

torchtyc check <paths>...        Shape-check files or directories
torchtyc run <file.py::entry>    Run an entry point on meta tensors
torchtyc trace <file.py::func>   Show the shapes flowing through one function
torchtyc watch <paths>...        Re-check on change
torchtyc lsp                     Language server on stdio
torchtyc lsp --tcp PORT          Language server on a TCP port
torchtyc rules                   List the diagnostic rules
torchtyc version                 Print the version

check takes --format full (default), concise, json, or github.

$ torchtyc trace model.py::Linear.forward
Linear.forward
  x       : (..., in_features)
  return -> float32[(..., out_features)]

dimension names are bound to distinct primes starting at 101

Constructing modules

To check a method, torchtyc needs an instance, and to build an instance it needs constructor arguments. It matches integer parameters of __init__ against the dimension names in the method's annotations:

def __init__(self, in_features: int, out_features: int) -> None: ...
def forward(self, x: Float[Tensor, "... in_features"]) -> ...

in_features is a dimension name, so it receives that dimension's prime. An int named dim or axis is passed as -1 unless an annotation names it as a dimension. Parameters with defaults are left alone. A parameter that is neither a dimension, a known type, nor defaulted produces an unresolved-arg warning and the function is skipped rather than guessed at.

Modules are built inside torch.device("meta"), and the initialisers in torch.nn.init are neutralised for the duration, since initial values cannot affect a shape.

Annotated attributes

Python does not check variable annotations at runtime, and neither does jaxtyping, which only reads function signatures. Since torchtyc has a constructed instance in hand anyway, it checks them too:

# attr.py
class Linear(nn.Module):
    def __init__(self, d_in: int, d_out: int) -> None:
        super().__init__()
        self.W: Float[nn.Parameter, "d_out d_in"] = nn.Parameter(torch.empty((d_in, d_out)))

attr.py:9:17: error[attribute-mismatch]
  `self.W`: annotated `d_out`, but the traced dimension is `d_in`
    Expected: (d_out, d_in)
    Got:      (d_in, d_out)
    9 | self.W: Float[nn.Parameter, "d_out d_in"] = nn.Parameter(torch.empty((d_in, d_out)))
  try:  Float[nn.Parameter, "d_in d_out"]
  hint: this dimension is `d_in`, so the annotation likely names the wrong axis

Running unannotated code

Most PyTorch code carries no shape annotations. torchtyc run checks it anyway, by running the script's real entry point with every tensor on the meta device:

torchtyc run tools/eval.py::main            # call main()
torchtyc run tools/eval.py                  # run the `if __name__ == "__main__":` block
torchtyc run tools/eval.py::main -- --n 6   # arguments for argparse

No weights load and no kernels run. from_pretrained builds the model from its config on meta, safetensors files are read from their header alone, and .to("cuda") stays on meta, so a multi-billion-parameter model runs a forward pass in a second on a machine without a GPU. A shape bug raises at your line:

tools/eval.py:36:5: error[run-error]
  `rearrange` raised EinopsError: Wrong shape: expected 2 dims. Received 3-dim tensor.
    Got:      logits (1, 113 (seq_len?), 250624 (vocab_size?))
    36 | rearrange(logits, "1 seq_len -> seq_len ")
  via: main (line 47) → run (line 36)
  hint: the pattern names 2 axes but the tensor has 3
run tools/eval.py::main: stopped at tools/eval.py:36 (run-error)

A name in parentheses is a guess: torchtyc prints one only when exactly one config key or einops axis has that size, and always prints the number.

A run stops, with run-stopped at info level, when the script needs a tensor's value (.item(), if loss < best, .tolist()), a file that is missing, a download, or a write. That is never reported as an error, and the status line says how far the run got.

torchtyc run accepts the same --format options (full, concise, json, github) and common flags (--python, --allow-effects, --severity) as check. For run, --timeout sets the per-entry run limit (overriding run-timeout), while for check it sets the trace-pass timeout.

The status line and JSON run_status describe the outcome of each entry:

• ran to completion: the entry finished without raising an unhandled exception;
• stopped: the run could not continue for an expected reason (such as a value read, missing input, blocked write or network, or non-zero SystemExit);
• timed out: the entry exceeded its per-entry run-timeout;
• crashed: the worker process terminated unexpectedly before the entry finished;
• not run: skipped before execution (such as a function requiring arguments with no default).

Exit codes match check:

• 1 when there is an unhandled run-error (or another error-level diagnostic);
• 2 when torchtyc itself failed (an invalid entry target, missing file, or worker process that died);
• 0 otherwise.

List entry points in pyproject.toml and torchtyc check runs the ones under the paths it checks:

[tool.torchtyc]
run = ["tools/eval.py::main", { entry = "tools/train.py", args = ["--steps", "1"] }]
path-map = { "/workspace" = "." }   # hard-coded container paths
run-timeout = 300                   # seconds per entry

The CI job tests torchtyc run against pinned transformers releases: 4.45.2, 4.57.6 and 5.20.0.

Suppressing

y = x.reshape(-1)  # torchtyc: ignore
y = x.reshape(-1)  # torchtyc: ignore[rank-mismatch]

A scoped ignore that never matches is itself reported, so suppressions do not rot.

Configuration

[tool.torchtyc]
python = ".venv/bin/python"   # interpreter that imports your code
severity = "warning"          # drop anything below this level
ignore = ["unused-dim"]
exclude = [".venv", "build", "dist", "__pycache__", ".git"]   # replaces the defaults
variadic-rank = 2             # how many axes `...` stands for
einops = true
timeout = 60.0
extra-paths = ["stubs"]       # prepended to the worker's PYTHONPATH
allow-effects = false         # let imported code write files, open sockets, and spawn processes
run = ["tools/eval.py::main"] # entry points `check` runs on meta
path-map = { "/workspace" = "." }   # rewrite container paths for loaders
run-timeout = 300.0           # seconds per run entry

Editors

Any LSP client works. Neovim, without a plugin:

vim.lsp.config.torchtyc = {
  cmd = { "torchtyc", "lsp" },
  filetypes = { "python" },
  root_markers = { "pyproject.toml", ".git" },
}
vim.lsp.enable("torchtyc")

The server publishes lint diagnostics immediately on every change, and traces after the buffer has been quiet for 0.7s, on open, and on save. It never imports your code on a keystroke.

Hover over a function to see the traced shapes. Inlay hints show the traced return next to each signature, and a code lens above it reports the traced return or the error count. Code actions offer to silence a rule or to adopt the shape that was actually traced.

CI

- run: uv run torchtyc check src/ --format github

--format github emits workflow commands, so each finding becomes an inline annotation on the pull request diff. check exits 1 when there are errors, and 2 when torchtyc itself could not do the job: the worker failed, or the paths matched no python file.

Rules

Rule	Level	Meaning
shape-mismatch	error	a traced shape disagrees with its annotation
rank-mismatch	error	a traced value has a different number of dimensions
dtype-mismatch	error	a traced dtype is outside the annotated dtype set
dim-inconsistent	error	one dimension name is bound to two different sizes
attribute-mismatch	error	an annotated attribute on self holds a different shape
not-a-tensor	error	an annotated tensor position received a non-tensor
tuple-arity	error	a tuple return has a different length than annotated
einops-pattern	error	an einops pattern disagrees with the tensors given to it
trace-error	error	the function raised while being traced
import-error	error	the module could not be imported
worker-error	error	the file could not be traced
run-error	error	the entry point raised while running on meta tensors
device-mismatch	warning	a traced value left the meta device
einops-unknown-axis	warning	an einops axis matches no input axis or keyword
uninstantiable	warning	a module's __init__ could not be called automatically
unresolved-arg	warning	a parameter has no annotation and no default
unsupported-annotation	warning	an annotation could not be parsed
run-unsupported	warning	the run reached something the meta runtime cannot stand in for
local-definition	info	a target inside a function body cannot be reached after import
anonymous-return	info	arguments are annotated but the return is not
missing-annotation	info	a public function has no jaxtyping annotation
trace-retried	info	a function traced only after widths were rescaled
unused-dim	info	a dimension name is used once, so it constrains nothing
suppression-unused	info	an ignore comment matched no diagnostic
run-stopped	info	the run could not continue for a reason that is not a bug

Limits

torchtyc imports your module, so module-level side effects run. Keep training loops behind if __name__ == "__main__":.

Classes and functions nested inside other classes, or written under a module-level if or try, are checked like any other. One inside a function body is not: after import there is no name to reach it by, so it reports local-definition rather than passing silently. A if TYPE_CHECKING: block is skipped without a diagnostic, because nothing in it exists at runtime.

It runs one concrete trace, not a proof. A function whose control flow depends on tensor values rather than shapes takes whichever branch the primes send it down. ... stands for a fixed number of axes, two by default, so code that behaves differently at other ranks needs variadic-rank or a second annotated wrapper. Every bare ... in one signature also stands for the same axes, which is narrower than jaxtyping: a loss taking two "... d" arguments traces with one batch shape, because that is what the annotation almost always means.

A dimension is a prime, and a prime does not divide, so code that splits an axis - head_dim = d_model // n_heads, then view(b, s, n_heads, head_dim) - gets a quotient that does not multiply back. A trace that fails this way is run again on widths that do divide, taken from the numbers the model itself writes down: a parameter's default, or a literal in the body of the constructor or the traced method. Multi-head attention passes on the second attempt and reports trace-retried at info level. A model that splits by a width written nowhere either of those can see, or by one above 256, still needs # torchtyc: ignore[trace-error].

The same retry covers a constructor that rejects odd widths, such as if d_k % 2: raise ValueError(...), and it runs for attribute checks as well as for forward. If the retry gets past the guard and hits an error, that error is the trace-error, and a trace-retried note gives the widths it needs. If the constructor instead refuses an argument torchtyc made up, with a raise or assert of its own or in a helper from the project, or fails the same way again, the result stays uninstantiable. A raise inside an except is judged by the error it caught, so a real bug that the constructor wraps and raises again is still a trace-error.

Runtime checking with jaxtyping and beartype remains worth having. torchtyc tells you the shapes are consistent for the sizes it chose; beartype tells you they were right for the batch you actually ran.

If your package installs one of those hooks over itself, remember that torchtyc calls your constructors with real integers, so your own runtime checks see them. A width annotated float that receives 101 is rejected by beartype, which does not apply the PEP 484 numeric tower unless asked, and it reaches you as a trace-error on the constructor. Every static checker accepts an int there, so the two disagree about your code rather than about torchtyc. Either annotate int, or turn the tower on for the whole package:

install_import_hook(
    modules="mypkg",
    typechecker="beartype.beartype(conf=beartype.BeartypeConf(is_pep484_tower=True))",
)

torchtyc run is one concrete execution, like the trace pass. Scripts built on generate() usually stop at the generate call, because sampling reads values. A value read made inside C++ (torch.allclose) is recognised by its message, so a bare except around it can hide the stop. path-map rewrites paths that loaders and open read; it does not rewrite os.path.exists, os.listdir or Path.glob, so a script that lists a mapped directory needs the real directory or a symlink. PEFT, DeepSpeed, FSDP and datasets are not supported, and end the run with run-stopped or run-unsupported. Transformers 4.x on meta needs the meta-aware torch.autocast the runtime installs (it is automatic), and uv run --with is how the CI job swaps versions. Device checks in user code (assert x.is_cuda, comparing x.device to "cuda") see the meta device and fail.

Security

torchtyc imports your modules in a worker subprocess that runs with your full privileges.

While it imports and traces your code, a guard blocks filesystem writes, outbound network calls (including DNS resolution), and process spawning. It catches accidental effects such as telemetry setup, dataset downloads (including those shelling out to git or curl), and checkpoint writes. It is not a security boundary: Python cannot enforce one in-process, and code written to get around it can. Set allow-effects = true if a model you trust needs to write, connect, or spawn processes.

Untrusted code still needs a real sandbox or a disposable container. Only enable torchtyc lsp in workspaces you trust.

Licence

Apache-2.0.