Coming from jaxtyping? Convert once, check forever¶
Already have Float[np.ndarray, "batch features"] everywhere? For new
code, prefer the native spelling — same feel, no middleman. For
existing files, a plugin cannot read jaxtyping's dim strings — they ride
inside Annotated metadata, which no mypy type-analyze hook ever sees
(probed every hook: only jaxtyping's internals fire, the string never
arrives). So the bridge is a one-time source rewrite instead of a hook:
numty_convert mechanically
translates the annotations, then gets out of the way:
# before
def forward(x: Float[np.ndarray, "batch inp"], ...) -> Float[np.ndarray, "batch out"]: ...
# after (import of jaxtyping dropped — nothing references it anymore)
BATCH = TypeVar('BATCH', bound=int)
INP = TypeVar('INP', bound=int)
def forward(x: FlexArray[tuple[BATCH, INP], np.dtype[Any]], ...) -> FlexArray[tuple[BATCH, OUT], np.dtype[Any]]: ...
Names become shared TypeVars — same name means same size,
case-sensitive like jaxtyping (B and b stay two variables) —
fixed sizes become Literals, ..., _ and unchecked _name become
fresh variables, *name becomes a named *Ts (with the
# type: ignore[type-var] such signatures need), exact dtypes (Float64,
Int32, …) keep their precision, and family heads (Float, Int, …)
keep the shape while the dtype stays Any. Only def parameter and
return annotations convert. The rest of the file is left byte for byte:
comments, pragmas and blank lines survive, and a file with nothing to
convert comes back unchanged. It catches real mistakes: mis-annotate the output
bias as "hid" instead of "out" and the converted file fails with
error: broadcast at trailing axis 0 (shapes (BATCH, OUT) vs (HID,)): OUT and HID are different names [numty-names]
error: call-site shape conflict for 'HID': 8 (argument 'w1', axis 1) vs 2 (argument 'b2', axis 0) [numty-callsite]
Two different names are two sizes. jaxtyping accepts "hid" and
"out" that happen to be equal at run time; numty's numty-names
code, on by default, reports a body that needs them to be equal (the
first error above, where the mistake is made). Where two names really
are one size, give them one name, or turn the code off with
disable_error_code = numty-names (see
Error codes).
Anything the static solver cannot express is left verbatim with a
warning, never silently changed: symbolic dims ("n+1"), broadcastable
#name dims (size 1 or the name — no static spelling), two variadics,
stacked subscriptions, non-numpy bases, aliased heads (Float as F),
string annotations, jaxtyping subscripts outside a signature (module
variables, isinstance), and names your code already binds (a
module-level BATCH = 5 or from config import BATCH would be rebound by
the emitted decl, so that annotation stays jaxtyping). A fully converted
file runs unchanged — the annotations evaluate — and no longer needs
jaxtyping installed. The round trip is pinned by
tests/test_convert.py:
goldens for every dim kind plus an end-to-end MLP that checks clean and a
seeded-bug twin that must fail.