Shapes and sizes

Tensor types

A tensor type is a shape, then an optional domain, then an optional dtype: [b n D] Signed bf16. Each axis of a shape is one of:

Whether an axis has a name matters to broadcasting and to over=, which find axes by name.

A type with a domain or dtype but no shape, as in x: Real or p: Unit, makes a shapeless parameter or return type, as in an elementwise function. In as and to, it keeps the value’s shape.

Size polynomials

Every size is a polynomial with integer coefficients over size symbols: dims, generic axes, unknown sizes and the variables the solver works out. 3*D, gh*gw and n+1 are polynomials. They are kept in a normal form, a sum of distinct products of symbols, so equal sizes are equal however they were written: 3*H*Dh and 3*Dh*H are the same.

Two sizes are compared after replacing each dim by its value. So with dim D = 768, dim H = 12 and dim Dh = D / H, the sizes 3*D and 3*H*Dh are both 2304 and are equal. Generic and unknown sizes stay symbolic: n equals n and nothing else, and n+n equals 2*n.

Arithmetic on sizes is checked: a size that doesn’t fit in 128 bits is E0414.

Division

a / b in a size is exact division. It is allowed when b is a single term that divides every term of a, symbolically or once dims are replaced by their values: 2*n / 2 is n, D / H is 64 with the dims above, and n*m / n is m. Anything else, such as D / 5 with D = 768 or n / 2, is E0402. Write the other size as a multiple instead: [2*n] and [n].

In the return type of an extern fn, division rounds down (Rounded division).

Dims

dim D = 768 declares a constant size. Its value is a size expression over numbers and other dims, evaluated in dependency order. A dim defined in terms of itself is E0304. A dim’s value must be a whole number of at least 1, and a division in it must be exact (E0305). A name in a dim’s value that isn’t a dim is E0301 or E0303.

In the editor, hovering a dim shows its value and docs, then where it’s used: the named values whose types have it in a size (parameters, values in function bodies, and fields of models and records), grouped by the function, model or record they’re in, each with its type, as `w_qkv: [D 3*D]`. The types show where the dim sits, which names alone wouldn’t. At most six groups of four values are listed, and the rest are counted.

Grouped sizes

A number touching a name, as in 16gh, is a grouped size: one axis of size 16 × gh. It is the group (gh 16), so in a parameter’s type it introduces gh, a generic axis that the rest of the signature and the body can use on its own, as in fn patchify(img: [b 3 16gh 16gw]) -> [b (gh gw) D], and a call solves for it: 16gh against 224 gives gh = 14, and against 225 is E0402. In a pattern it is gh outer and 16 inner, the layout of (gh 16); the other layout is written (16 gh). The number is an integer and the name an axis name; a dim’s name, as in (gh P), needs the group’s parentheses.

16 gh, with a space, is two axes, 16 and gh. Where a grouped size was probably meant, check and the editor give the hint H2201: for a number greater than 1 directly followed by a name in a function’s signature, when the name isn’t a dim, is an axis of its own nowhere else in the signature, and is part of a group somewhere in it, but not a grouped size, as gh is in fn f(img: [b 3 16 gh]) -> [b (gh gw) D]. A name that is an axis of its own elsewhere, as n in fn f(x: [2 n]) -> [n], gets no hint, and neither does one written as a grouped size elsewhere, as gw in fn rows(img: [b 16gh 16gw]) -> [b gh 16 gw 16].

Shape aliases

dim Images = [b 3 16gh 16gw] declares a shape alias: a name for a list of axes. It can be a type’s whole shape, x: Images, with an optional domain and dtype after it, as in x: Images Unit f32, or among another shape’s axes, [Images D], which is [b 3 16gh 16gw D]. The alias’s axes replace its name after parsing, so the names in it are looked up where it’s used: in fn f(x: Images) -> [b D], b is f’s generic axis, the same in the return type, and each function that uses the alias has its own b. An alias may use other aliases, declared before or after it, but not itself, directly or through others (E2201). An alias stands for several axes, so it can’t be part of one axis, inside a group or arithmetic (E2202), and it isn’t a value or a size (E0303). In patterns, a name is always an axis name, never an alias. The editor shows what an alias stands for on hover, and goes to its declaration from a use.

Another file’s alias is written with that file’s name, as its other declarations are (Importing other files): after import vision, x: vision.Images, x: vision.Images Unit f32 and [vision.Images D] work as they would in vision.ss. The dims in the alias are named as the importing file names them, so dim Tokens = [b n D] in vision.ss is [b n vision.D] here; its other names, such as b and n, are still the using signature’s own. Aliases are replaced file by file once every import is read, so an alias may also use another file’s alias, as dim Pair = [vision.Images 2].

Unknown sizes

An unknown size is a fixed size that the checker can’t see. It is written ?k and shown as ?k. It equals only itself, so [?k D] + [196 D] is an error (E0401) until an as gives the size: rows as [196 D]. Where it matters, as when slicing, an unknown size is assumed to be at least 0, and other sizes at least 1.

Unknown sizes come from:

• indexing with a mask: x[sum(x, over=D) != 0] is [?k D];
• a slice that may not fit its axis, or whose bounds are only known at run time (Indexing);
• topk with a k that is a run-time value, and a pattern keyword size that is a run-time value;
• ?k in a function’s return type: inside the function, ?k is whatever size the body returns, and each call gives a new unknown size;
• ?k in the type of an as or to: a new unknown size each time the expression is checked.

Within one type, the same ?k is the same size. Unknown sizes are numbered per function: the second one named k is shown as ?k2. ?k can’t appear in a parameter’s type (E0306): the caller chooses parameters’ sizes, so write k.