Signature style — shape comments as the real type system#

A house convention for T3Toolbox function signatures. Read this once; apply the principle everywhere, the mechanics within reason. Reference exemplar: t3toolbox/backend/probing.py. Shape vocabulary (C/W/K blocks, elm_shape, single-axis letters): docs/batching_and_stacking.md.

The principle (this is the part that matters)#

In array code the nominal type is nearly useless — the shape is the type. Every argument is NDArray (or Union[Sequence[NDArray], NDArray]), so Python’s annotations tell a caller almost nothing: not the rank, not which axes are batch vs core, not the length of a core tuple, not which bond matches which. We recover that lost contract with a trailing comment after each argument and each return element. The comment is not prose decoration — it is the type annotation that the language cannot express.

Two goals, both about the reader:

  1. Fill the gap in Python’s type system for array/tensor code.

  2. Make a caller understand the arguments at a glance — jump-to-definition should answer “what shapes does this need, and what comes back?” without reading the docstring or the body. That is where people actually look.

If an edit serves those two goals, it’s in the spirit of the rule. If a mechanical detail ever works against them, the goals win — see “Within reason” below.

Annotations vs comments: express what the type system can, comment the rest#

The comment-as-type technique is a patch for a specific deficiency — array shape. It is not a license to write ordinary types as comments. The division of labor:

  • Whatever Python can express, put in the annotation. A plain int | Sequence[int] | None, Optional[float], bool, a callable — fully expressible, so annotate them (real, checkable, IDE-visible; this is a library for outside users whose tooling consumes annotations). Be correct-and-explicit: include None in the union / use Optional.

  • Whatever it can’t, put in the comment. Array shapes (NDArray is content-free), the length of a sequence (Sequence[int] can’t say len=d+1), cross-argument constraints (1 <= max_rank <= min(N,M)), and semantics (requires unstacked).

The two coexist on one argument, exactly like array args do (a weak NDArray / Union[...] annotation plus a load-bearing shape comment). Example:

max_tt_ranks:  typ.Union[int, typ.Sequence[int], None] = None,  # scalar caps all, or len=d+1
rtol:          typ.Optional[float] = None,                      # requires unstacked

Why not comment-only? Because annotating the expressible part keeps the presence of a comment meaningful: a # then signals “there is a contract here the type system cannot capture.” If comments also carried ordinary types, you couldn’t distinguish load-bearing shape comments from types-written-as-comments. (It also keeps the annotation column uniform — every arg carries its Python type; the comment carries the rest.)

The mechanics#

def compute_mu(
        left_tt_cores:  typ.Union[typ.Sequence[NDArray], NDArray],  # len=d-1, elm_shape=C+(rLi,nUi,rL(i+1))
        xis:            typ.Union[typ.Sequence[NDArray], NDArray],  # len=d, elm_shape=W+C+(nUi,)
) -> typ.Union[
    typ.Sequence[NDArray],  # mus. len=d, elm_shape=W+C+(rLi,)
    NDArray,
]:

  1. def name( on its own line; one argument per line as name: type, # contract.

  2. The closing ) -> ...: carries the return shape(s), each return element on its own line with its own trailing contract — expand return tuples one-element-per-line even when they would fit on one line.

  3. Three vertically-aligned columns: argument name · type annotation · # contract. The block reads as a table; a single-arg change touches a single line.

  4. The contract uses the shared shape vocabulary — the same letters as the body-local suffixes (mu_WCa, B0_b_j_c) and the grouped-contraction subscripts ('WCa,Caib,WCi->WCb'). One vocabulary across signature, body, and helper names: learn it once, read it everywhere.

Micro-grammar (keep the table cells uniform)#

  • Arrays: # len=<n>, elm_shape=<...> — len first (for core tuples), then elm_shape, comma separator. Use shape=<...> when there is no tuple wrapper (a bare array).

  • Non-array / scalar args: a short role or constraint note, e.g. # 1 <= max_rank <= min(N,M), # 'left' | 'right', # requires unstacked. If a scalar’s role is obvious from its name, a note may add nothing — then keep it short or omit (see below).

  • Placement / dtype, when load-bearing: NDArray says nothing about dtype or host-vs-device either — so when those are part of the contract (not just shape), put them in the comment too. The motivating case is the uniform masks, tagged HOST bool, static alongside their shape. The HOST token means “must be a host (numpy) array, never jax/traced — it is static structure, runtime-enforced by the mask guard” (under jit a traced mask breaks; see uniform_pytree_composition.md). No tag = the default: dtype-agnostic and xnp/backend-polymorphic (numpy or jax). Spell a dtype (bool, int) only when it actually matters. This keeps the old “a bare np. is a backend-agnosticism tell” intuition usable: a HOST arg is meant to be numpy.

  • Prefer the comma form over a period (# len=d, elm_shape=..., not # len=d. elm_shape=...).

Aligning the comment column: one block, or grouped?#

Names and types are always column-aligned — each is an O(1) jump for the eye. The only open question is the comment column. The cost of forcing every comment into a single column is the type-length spread: the gap an argument gets is (longest type+default) − (its type+default). So:

  • Small spread → one aligned comment column. Cheap, tidy, every comment O(1). The common case (e.g. probing.py, where every arg is the same array type — the column falls out for free).

  • Large spread → group, and align within each group. Split the args into runs of similar-length types (which usually coincide with roles: operands / rank caps / tolerances / flags), put a blank line between groups, and align the comment column within each group. Group by type-length, not role — length is what sets the gap; role is only the usual proxy. When they disagree (e.g. a short str selector sitting among long Union caps), an argument whose type-length matches no group gets its own blank-delimited group rather than being force-aligned into a column it doesn’t fit — otherwise you reintroduce the very gap grouping exists to remove. The test is the reader’s eye, not a column count: split an arg out when force-aligning it would open a gap big enough to lose the row (a few spaces is fine; a tab-plus usually is not). The comment column resets per group; names/types stay aligned across the whole signature (still O(1)) — only the comment column is per-group. This keeps every comment O(1) within its group, keeps lines short, and makes the grouping visible. It is strictly better than one column with 30-space gaps, and far better than the staircase (each comment 2 spaces after its own type, no cross-row alignment): the staircase minimizes gaps but makes the comment column unpredictable — O(n) to find — which is the worst to scan. Don’t use the staircase.

Rule of thumb: group when the args span clearly different type-lengths (a long array/Union next to a bare float/bool) and there are enough rows for a staircase to bite (≳4). For two or three args a single column is already fine. (t3m, frontend and backend, is the worked example: groups operands / max_* caps / rtol,atol / oversample, blank-line-separated.)

Within reason (principle over dogma)#

This is a convention, not a contract enforced by anything.

The iron imperative for every comment is that it be useful to a user of the code or a maintainer of the library. A comment that only repeats what the signature already says is not neutral — it is a liability: a second, unverified copy of the contract, free to drift out of sync with the first. The shape comment exists for one reason: to carry type information the language cannot express (shapes, sequence lengths, which bond matches which). Where that information is already expressed — by an annotation, or by the name itself — the comment has no job to do.

Apply the convention where it earns its keep; handle the exceptions case by case. These rules cannot anticipate every case: the deviations below are illustrative, not exhaustive, and it matters more to follow their spirit than their wording. Reasonable deviations:

  • Trivial scalars (use_jax: bool, an obvious axis: int) may need no comment, or only a brief role note. Don’t manufacture filler to fill a column.

  • Self-evident or non-array types (a callable, a plain Python int count, an enum-like string) need only a short note, or none.

  • Pathologically long types/contracts (a deep nested Union of two full representations, as in t3_probe’s x) may wrap across several lines; perfect single-line alignment is impossible there, so keep it readable rather than rigidly tabular.

  • One-off / throwaway internal closures (_func inside an xmap) don’t need the full treatment.

  • Grouped-contraction call sites (contractions.contract) — an explicit exception. There the subscripts string already is the shape contract ('WKCa,Caib,WCi->WKCb' says operand 1 is W+K+C+(a,), operand 2 is C+(a,i,b)), so a comment restating it tells the reader nothing the call has not already told them — exactly the liability above. The supplement keywords state their own contract by name (len_C= = the number of C axes). Prose describing an operand’s role at one call site (“mu jet”, “probe vector w”) belongs in the calling function, never inside contractions.py: the interpreter is deliberately caller-agnostic, a mechanical workaround for einsum’s lack of grouped indices.

The two costs we accept knowingly: the comments are unverified (they can drift — tests and review are the only guard, so update them when shapes change), and the alignment is hand-maintained (no autoformatter; one long name re-flows the block). Both are the price of using comments as types; we pay it on purpose.

Scope notes#

  • These trailing comments are the source-reader / jump-to-def surface. The numpydoc Parameters/Returns docstring is the rendered-docs surface. They may both carry shapes today; if so, keep them consistent. (Whether docstrings should stop repeating shapes is a separate, undecided question — don’t conflate it with signature cleanup.)

  • Applies to the frontend (TuckerTensorTrain & friends) too, not only the backend — and to the uniform layer (whose HOST bool, static mask tags are the motivating case above).

  • The parked weighted layer (wt3_*) is not held to this until its post-1.0 revival.

  • The reader-side half of this convention (the decode table users need) is ../reading_signatures.md.