ShapeScript language reference

This is the reference for the ShapeScript language as the compiler in this repository implements it. The README is a tour; this document is the precise version. When the two differ, this document describes what the compiler does, and the difference is a bug in one of them.

Each diagnostic code has a section under Diagnostics. The language server links every diagnostic to its section, and shapescript explain E0401 prints it.

Contents:

• Lexical structure
• Grammar
• Names and scoping
• Shapes and sizes
• Unification and the size solver
• Functions and calls
• Broadcasting by name
• Pipes
• Pattern operations
• Indexing
• Built-in operations
• Assertions with as
• Value domains
• Conversions with to
• Dtypes
• Extern functions and the prelude
• Host scalars
• Control flow
• Iterator chains
• Models
• Changing models with mut
• Gradients
• Optimizers
• Modes and randomness
• Data and datasets
• Logs and sinks
• Checkpoints
• Modules and generic functions
• Training loops
• Devices
• Several devices
• Fixed shapes and padding
• Python interop
• Runtime checks
• Export
• Notebooks
• Running programs
• Completion and signature help
• Debugging
• Not yet implemented
• Diagnostics

Lexical structure

Source text

A source file is UTF-8. A byte order mark is ignored. Spaces, tabs, form feeds and carriage returns between tokens are ignored, except that indentation is significant and may only use spaces.

Comments and doc comments

# starts a comment that runs to the end of the line. ## starts a doc comment.

A doc comment documents the declaration on the line directly below it. Consecutive lines that hold only a ## comment form one doc comment; each line’s ## and one space after it are removed. A doc comment documents a dim, domain, fn, extern fn, import python, model, or a model’s field, and the editor shows it when you hover the declared name. A doc comment that isn’t directly above a declaration, or that follows code on its line, is a warning (W0201).

Math in comments

Any comment, # or ##, can hold LaTeX math: $...$ inline, or $$...$$ displayed. The rules for finding math are Pandoc’s:

• An opening $ can’t be followed by whitespace, and a closing $ can’t come after whitespace or be followed by a digit, so costs $3 and $4 is text.
• \$ is a literal dollar sign, a run of three or more $ is text, and nothing inside a Markdown code span (`$x$`) is math.
• Inline math stays on one line. Display math may continue over consecutive comment lines of the same kind (all # or all ##, each with nothing but the comment on it), but not past a blank comment line, and can’t be empty. A comment after code is a block of its own.

Math accepts the LaTeX subset listed in the README’s Math in comments table. Anything else is shown as written, with a hint: H0601 for a command or environment outside the subset, and H0602 for malformed math. Math never causes an error; the hints come from shapescript check and the editor, not from build or the Python import hook.

Tokens

Keywords: dim, domain, fn, extern, return, import, as, to, and, or, inf, true, false, for, in, if, else.

Contextual words are identifiers with a special meaning in one place: python after import; over after a domain in a type; devices at the start of a top-level line (in a notebook cell, only when a name follows it; Devices); model at the start of a top-level line when a name follows it; export at the start of a top-level line when fn follows it (Export); and param and buffer at the start of a line in a model’s block (Models). Elsewhere they are ordinary names, so model = Mlp.init(seed=0) in a notebook cell binds a value named model. int and float are ordinary names too, which mean host scalars in a type.

inf is infinity; -inf is unary minus applied to it, so x-inf means x - inf. true and false are the two bools. Any other character is E0101.

Lines and indentation

The layout pass turns line breaks and indentation into NEWLINE, INDENT and DEDENT tokens:

• Inside ( ) or [ ], line breaks and indentation are ignored, so a bracketed expression can span lines.
• Blank lines and lines with only a comment are ignored.
• A line whose first token is |> continues the line before it. No NEWLINE comes between them, and its indentation isn’t checked.
• Otherwise each line ends with a NEWLINE. A line indented more than the current block opens a block (INDENT); a line indented less closes blocks (one DEDENT each) until its indentation matches an enclosing block’s. If it matches none, that’s E0104.
• Indentation is counted in spaces. A tab in indentation is E0103.

An unmatched ) or ] is E0105, and a bracket that is never closed is E0106.

Grammar

Notation

The grammar below is EBNF: { x } is zero or more, [ x ] is optional, | separates alternatives, and quoted text is a token. NEWLINE, INDENT and DEDENT come from the layout pass. ident, int, float, string and unknown_size are the tokens of those names. The parser is hand-written recursive descent; where the grammar can’t say something, the text after it does.

The parser never stops at an error: it reports the problem, skips to the end of the line (and past any block indented under it), and continues, so one file can report several syntax errors. A misplaced token is E0201.

Files and items

file        = { NEWLINE | devices_decl | item } ;
item        = dim_decl | domain_decl | fn_decl | extern_decl | import_decl | model_decl
            | data_decl | logs_decl ;

devices_decl  = "devices" ident [ "[" { int } "]" ] NEWLINE
                [ INDENT { device_option NEWLINE } DEDENT ] ;
device_option = ident { token } ;

dim_decl    = "dim" ident "=" size_expr NEWLINE
            | "dim" ident "=" shape NEWLINE ;        (* a shape alias *)
domain_decl = "domain" ident "=" interval NEWLINE ;
interval    = ( "[" | "(" ) expr "," expr ( "]" | ")" ) ;
import_decl = "import" "python" ident { "." ident } NEWLINE ;

fn_decl     = [ "export" ] "fn" ident "(" [ params ] ")" [ "->" type ] fn_body ;
fn_body     = ":" NEWLINE INDENT { stmt } DEDENT
            | "=" expr NEWLINE ;
extern_decl = "extern" "fn" ident "(" [ params ] ")" [ "->" type ] "=" string NEWLINE ;

params      = param { "," param } [ "," ] ;
param       = [ "mut" ] ident [ ":" type ] [ "=" expr ] ;

model_decl  = "model" ident ":" NEWLINE INDENT { field } DEDENT ;
field       = ( "param" | "buffer" ) ident ":" type [ "=" expr ] NEWLINE ;

data_decl   = "data" ident ":" NEWLINE INDENT { data_field } DEDENT ;
data_field  = ident ":" type NEWLINE ;

logs_decl   = "logs" ident ":" NEWLINE INDENT { metric } DEDENT ;
metric      = ident ":" type ident NEWLINE ;

A file holds only declarations; a statement at the top level is E0207, except in a notebook cell. A declaration inside a function is E0201. In an interval, [ and ] include the end and ( and ) exclude it, so (0, 1] is half-open.

A file has at most one devices_decl, before every item (Devices). Its ident is the device, such as cpu or cuda, and its ints are device numbers. A model’s block holds only fields (E0201); a field’s = expr is its initializer, which the parser allows to be missing and the checker doesn’t (Models). A data record’s block holds only fields, each a name and a type (E0201; Data and datasets). A logs record’s block holds only metrics (E0201). A metric’s last ident says how its values combine; it is read as part of the type first, so in loss: [] mean the type is [], and in loss: [] mean last it is [] mean, with mean as a domain (Logs and sinks). model, data and logs are keywords only at the start of a top-level line and followed by a name, so they can still name values. export is a keyword only before fn, and marks the function for export; export extern fn isn’t a declaration (E0201).

Statements

stmt  = "return" [ expr ] NEWLINE
      | ident { "," ident } "=" expr NEWLINE
      | ident "." ident "=" expr NEWLINE      (* changes a model's field *)
      | "for" ident "in" expr ".." expr suite
      | "for" ident "in" expr suite
      | "if" expr suite [ "else" suite ]
      | expr NEWLINE ;
suite = ":" NEWLINE INDENT { stmt } DEDENT ;

The second form binds one name, or one name per result of split, topk or value_and_grad (E0411). A statement is a binding when its first two tokens are an identifier and = or ,. The third form, as in m.w1 = e, changes a field of a model (Changing models with mut).

a..b is part of a for header, not an expression, so a range can’t be bound to a name or passed to a function. else: goes on the line after the if block, indented like the if; an else anywhere else is E0201. There is no elif, while, break or continue: nest an if in the else block instead of elif. What loops and conditions mean is in Control flow.

Expressions

expr      = range { "|>" stage } | lambda ;
range     = or_expr [ ".." or_expr ] ;          (* a range only starts an iterator chain *)
lambda    = params "=>" expr ;                  (* only the function of map or filter *)
params    = ( ident | "_" ) | "(" ( ident | "_" ) { "," ( ident | "_" ) } ")" ;
stage     = "to" type | or_expr ;               (* see Pipes for which stages are valid *)
or_expr   = and_expr { "or" and_expr } ;
and_expr  = cmp_expr { "and" cmp_expr } ;
cmp_expr  = sum [ cmp_op sum ] ;
cmp_op    = "==" | "!=" | "<" | "<=" | ">" | ">=" ;
sum       = product { ( "+" | "-" ) product } ;
product   = unary { ( "*" | "/" | "@" ) unary } ;
unary     = ( "-" | "!" ) unary | cast ;
cast      = postfix { ( "as" | "to" ) type } ;
postfix   = atom { call_args | index_args | "." ident } ;
atom      = ident | int | float | string | "inf" | "true" | "false" | "_" | "(" expr ")" ;

There are no list literals, so [x, y] is E0201, and [e for i in n], the comprehension that iterator chains replaced, is E3201, whose help gives the chain to write. A lambda, i => e or (a, b) => e, may only be the function of a map or filter stage (E3202), and its body runs to the end of the argument: map(v => v * 2.0, over=n). A range a..b may only start an iterator chain, 0..n |> map(…), or be what a for loop loops over (E3203). In map(…) and filter(…), to T alone is the function v => v to T.

A stage is parsed like an or_expr, but must be a name, a member access such as torch.exp, or a call whose callee is one of those; anything else is E0204. A _ may only be an argument of a call that is a pipe stage, at most once per stage (E0205). A pattern (an argument containing ->) may only be an argument of rearrange, reduce, repeat, einsum or split (E0203); a line with a -> outside brackets where an expression should be is reported the same way.

Operator precedence

From loosest to tightest. Binary operators are left-associative.

So x @ w |> gelu means (x @ w) |> gelu, -x as T means -(x as T), and a as [n D] + b means (a as [n D]) + b. A type ends where something that can’t continue it appears, so x as [n D] + b is unambiguous. To apply as or to to a whole pipeline, use parentheses: (x |> f) as T; x |> f as T is E0204.

Calls and arguments

call_args = "(" [ arg { "," arg } [ "," ] ] ")" ;
arg       = ident "=" arg_value            (* keyword argument *)
          | pattern                        (* an argument with "->" outside brackets *)
          | expr ":" pattern_many          (* an argument with ":" outside brackets *)
          | arg_value ;
arg_value = [ "mut" ] expr ;               (* "mut" only before a name *)

mut and param, buffer and model are contextual keywords. mut is one before a parameter’s name and at the start of an argument followed by a name, as in step(mut m, x); elsewhere it is an ordinary name.

The parser looks ahead to the next , or ) outside brackets to tell the forms apart: a -> makes the argument a pattern, and a : makes it an operand followed by a pattern with one input side per operand, as in einsum(q, k: b H n Dh, b H m Dh -> b H n m), where q and k are the operands.

Patterns

pattern      = side "->" side ;
pattern_many = side { "," side } "->" side ;
side         = { p_axis } ;
p_axis       = ident | int | grouped_size | "(" { p_axis } ")" ;

A pattern is plain syntax, not a string. 16gh in a pattern is the group (gh 16). What patterns mean is in Pattern operations.

Indexing and slices

index_args = "[" [ index { "," index } [ "," ] ] "]" ;
index      = [ expr ] ":" [ expr ] [ ":" [ expr ] ]     (* slice lo:hi:step *)
           | expr { expr } ;                             (* one expression, or several *)

Several expressions separated only by spaces, as in causal[n m], are only allowed for causal (Indexing).

Types

type      = shape [ domain ] [ dtype ]
          | alias [ domain ] [ dtype ]      (* a shape alias as the whole shape *)
          | domain [ dtype ]
          | dtype ;
shape     = "[" { axis } "]" ;
axis      = size_expr
          | "(" { axis } ")"                (* a group *)
          | grouped_size                    (* 16gh, the group (gh 16) *)
          | unknown_size                    (* ?k *)
          | "..." ;                         (* other axes; extern fn only *)
alias     = ident ;                         (* the name of a shape alias *)
domain    = ident [ "[" size_expr "]" ] [ "over" ident ] ;
dtype     = "bool" | "u8" | "i8" | "i16" | "i32" | "i64" | "f16" | "bf16" | "f32" | "f64" ;

An identifier after the shape is a dtype if it is one of the dtype names, and a domain otherwise. [] is a scalar’s shape. A group holding a single size may be followed by arithmetic, so (n+1)*2 is one axis. A grouped_size can’t: 16gh*2 is E2203, and so is a grouped size anywhere but an axis, as in an expression or a dim’s value. A type starts with an alias when its first name is declared as a shape alias anywhere in the file; a shape alias among a shape’s axes is an axis of the form size_expr, and both are replaced by the alias’s axes after parsing.

Some types are written as a single name, or a name with a model or a record, and aren’t tensor types. A type that is only int, float or bool is a host scalar: flag: bool is a Python bool, while keep: [n] bool is a mask. A type that is only a model’s name, as in m: Mlp, is that model, Grad[Mlp] is a gradient with respect to one, and AdamW[Mlp] and SGD[Mlp] are optimizer state for one. A data record’s name, with or without a shape before it, as in x: Digit or batch: [b] Digit, is one example or a batch of them, and Dataset[Digit] is a dataset (Data and datasets). A type that is only a logs record’s name, as in -> StepLogs, is that record, LogSink is a sink, Mode is a mode, and Schedule is a learning-rate schedule.

Axes are separated by spaces, but an operator joins its operands into one axis, so [b 3 * D] has two axes, b and 3*D, and [n - 1] has one.

Size expressions

size_expr = size_term { ( "+" | "-" ) size_term } ;
size_term = size_atom { ( "*" | "/" ) size_atom } ;
size_atom = int | ident | "(" size_expr ")" ;

What sizes mean is in Shapes and sizes.

Notebook cells

A %%shapescript cell is parsed like a file, except that statements, including for and if, may also appear outside any fn, and are checked in order. return outside a function is E0201, even inside a loop. A line starting with devices declares devices only when a name follows it, so devices = ... binds a value. See Notebooks.

Names and scoping

Module scope

The module scope holds the file’s dims, shape aliases, domains, fns, extern fns, models, data and logs records, the other .ss files that import binds (Importing other files), and the Python modules that import python binds. import python torch.nn binds torch, as in Python, and importing the same module twice is fine. Declarations can be used anywhere in the file, before or after they appear, and dims may refer to each other in any order.

A module-scope name can be declared once (E0302), except functions that each take a different model first (Overloaded functions). A user domain can’t have the name of a built-in domain or a dtype (E0503).

Below the module scope is the prelude, the extern fn stubs every file sees (The prelude), and below that the built-in operations (Built-in operations). A declaration in the file hides a prelude stub or built-in of the same name.

Function scope

A function body starts with its parameters, as values, and the axis names its parameter types introduce. A statement x = e binds x for the rest of the body; binding the same name again replaces it for the code after, and the new value may have a different type. Inside a for or if block the rules are stricter, and what the block binds is its own (Blocks and scope). A function’s names are not visible in other functions.

A name used as a value is looked up in this order: the function’s values and axis names, the module scope, causal, the built-ins, then (in a notebook) Python variables. An unknown name is E0301, with a suggestion when a name with a similar spelling exists.

Using a name as the wrong kind of thing is E0303: calling a value or a dim, using a function or domain as a value, using a value as a size, or using anything but a domain as a domain. A dim or an axis name used as a value is its size, a Python int.

Axis names

An axis name is a name for a size. Axis names are introduced in four places:

1. Parameter types. A name in a parameter’s type that isn’t in scope is a new generic axis, as b and n in fn f(x: [b n D]). The caller chooses its size (Generic calls). Later parameters can use it. A name in an Index[V] domain of a parameter introduces V the same way.
2. Patterns. A name in a pattern that isn’t in scope is introduced with the size of the axis it binds to (Pattern operations).
3. stack(..., over=k), when k is a new name.
4. Keyword sizes in patterns, as h=8 in rearrange(x, b n (h d) -> b h n d, h=8).

Once introduced, an axis name is in scope for the rest of the function, or of the for or if block that introduced it: in later types, patterns, over= arguments, causal[n m], and expressions such as sqrt(Dh).

Return types and the types in as and to can only use sizes already in scope; any other name is E0301 (“unknown size”). For a size the function itself chooses, a return type writes ?k (Unknown sizes).

A dim is not an axis name, but it can name an axis: in [b n D] the third axis is named D and its size is the dim’s value. A pattern or type that uses a dim’s name requires that size.

Name clashes

Inside a function, a value can’t take the name of an axis in scope or of a dim, and a parameter can’t take the name of an axis an earlier parameter introduced (E0303). A value may take the name of a function, which it then hides for the rest of the body. Outside any fn in a notebook cell, a value can’t take the name of any declaration, since it would replace the declaration for every cell that uses it.

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.

Unification and the size solver

Size symbols

Each size symbol is one of four kinds:

A function’s generic axes are rigid in its body: nothing is assumed about them, so n doesn’t unify with 196, and two different generics never unify with each other.

Solving

The checker states facts about sizes as equations, lhs = rhs, where lhs may contain solvable variables. Equations come from:

• each axis of each argument at a call, against its parameter’s axis (Generic calls);
• each axis of a returned value, against the return type;
• a pattern group such as (gh p1), against the axis it binds to (Pattern operations);
• the Index[V] sizes of arguments, after their shapes.

The solver substitutes dims by their values, then repeatedly looks for an equation with exactly one unsolved variable that appears in it linearly (not squared), and solves it: from var * a + b = 0, var = -b / a, which must be exact. It goes round until no equation makes progress. So linear equations and products with one unsolved factor are solved: (gh 16) against 224 gives gh = 14, and (gh p1) with p1 known gives gh.

Solver failures

Each equation fails at most once, with one of four codes:

A message says which axis failed and why, as in “axis 1: ?k is not known to equal 196” or “(p1 p2 c) = 16*16*3 = 768, not D = 512”.

Functions and calls

Function declarations

fn name(params) -> T followed by : and an indented block, or by = expr for a one-expression body. export fn declares a function the same way, and also checks that it can be exported without Python. A function can call functions declared later in the file. A function without a return type can’t call itself, directly or through others (E0306), since its return type is inferred from its body.

The return type is optional. Without it, a function’s return type is the type of the first value its body returns, checked when a call needs it. With it, every return is checked against it (E0401, E0403, E0501, E0502), and a block body that declares a return type must return a value on every path (E0307): a loop may run no times, and an if without an else may not run, so neither returns for certain. A bare return in a function that returns a tensor is E0307 too.

Parameters

A parameter has a type, a default, or both (E0306):

• x: [b n D]: a tensor parameter. Its shape introduces axis names.
• x: Real, p: Unit: a shapeless parameter. It takes a tensor of any shape, and of any dtype unless its type names one.
• n: int, rate: float, flag: bool: a host scalar parameter.
• m: Mlp: a parameter whose type is a model.
• beta = 1.5: a parameter with only a default takes any number, of the default’s kind (int or float).
• stride: Size[S] = 1, over: Axis[D]: only in an extern fn (Axis and Size parameters). In a fn they are E0306.

A parameter’s name can be declared once (E0302).

Defaults. A parameter with a default may be left out of a call. In a fn, the generated Python works a default out once, where the function is defined, so a default is a number, true or false, possibly computed from numbers and dims, as in eps: Pos = 1 / D. It can’t use the function’s parameters or axis names, isn’t a tensor or a Python value, and a parameter with a shape or a model’s type can’t have one (E0306). A default with a type is checked like an argument for it: steps: int = 2.5 is E0502, and eps: Pos = 0 is E0501. The defaults an extern fn allows are in Extern declarations and Axis and Size parameters.

Matching arguments

Positional arguments fill parameters in order, and keyword arguments (name=value) fill the parameter of that name. Too many arguments, a keyword that isn’t a parameter, a parameter given twice, and a missing argument for a parameter without a default are each E0406. A string where a tensor is expected is E0412.

Generic calls

At each call, the callee’s generic axes become fresh solvable variables, and each argument’s axes are matched to its parameter’s axes by position: the argument must have as many axes as the parameter (E0403), and each axis gives one equation. Names don’t matter here: a [t D] argument fits a [n D] parameter, with n = t. When the equations fail, the call is reported once per argument, as “f can’t take this as x”, with a note for each failing axis.

The result is the return type with every generic replaced by its solution. A result axis keeps a name when the name means something to the caller: a dim keeps its name, and a generic takes the name of the argument axis it was matched to, if that had one. So f([t D]) with fn f(x: [n D]) -> [n D] gives [t D], and f([8 D]) gives [8 D]. Each ?k in the return type becomes a new unknown size, so two calls to the same function give two different unknown sizes.

After the shapes, the Index[V] sizes of arguments are matched against their parameters’, which can determine generics that no shape does: fn f(ids: [b] Index[V]) -> [b V] works out V from its argument. A different size is E0507. Then each argument’s dtype and domain must fit its parameter’s (Widening, Dtypes).

Elementwise functions

A function whose parameters all have no shape, and whose return type, if any, has no shape, is elementwise: it applies to tensors of any shape. Its arguments broadcast together, and the result has the broadcast shape.

## Swish with a fixed slope.
fn swish(x: Real, beta = 1.5) -> Real = x * sigmoid(beta * x)

The declared return type is an upper bound on the result’s range. A call’s range is the narrowest of three: the declared one, the one the checker found for the body with the parameters’ domains, and the one an interval interpreter gives for the body with this call’s argument ranges (see Range inference). So fn signed(x: Real) -> Real = sigmoid(x) * 2 - 1 gives Signed values. A declared domain that isn’t a range, such as Mask, is kept; without ->, so is the body’s. The result has the arguments’ dtype unless the return type names one.

Activation curves on hover

In the editor, hovering a built-in activation (relu, leaky_relu, elu, gelu, silu, sigmoid, tanh, softplus, log_sigmoid) or a function of one number of your own shows its formula, its curve on $[-4, 4]$ (or the part of it in the input’s domain, $[0, 1]$ for Unit) with its derivative dashed, and the range of its values. A function of one number is:

• an elementwise function, plotted against its first parameter; or
• a function with exactly one tensor parameter, of any shape, generic or not (x: [b D], x: [b n h], p: [b n] Unit), wherever it is among the parameters, whose other parameters are all host scalars (k: float, beta = 1.5, smooth: bool = true), and whose result has the input’s axes, such as fn act(x: [b D]) -> [b D] = x * sigmoid(1.702 * x); or
• an extern fn of one tensor parameter for a Python function whose math the checker knows (torch.sigmoid, torch.nn.functional.gelu, silu, softplus, logsigmoid and the other entries of its table). The language server doesn’t run Python, so an extern fn for anything else, such as torch.nn.functional.mish, has no curve.

Its body must be elementwise in the tensor: every element of the result depends only on the same element of the input. That holds for a body built from numbers, inf, true and false, dims with a value, the parameters, values defined with =, unary minus, + - * /, comparisons, and, or, !, pipes, the elementwise built-ins (abs, square, pow, exp, clamp, where, maximum, the activations, …), calls to other functions of one number, and if/else whose condition uses only host scalars. Anything else isn’t: reductions (sum, mean), operations along an axis (softmax, layer_norm), @, indexing (x[0]), patterns (rearrange), as, to, loops, Python calls, an if on the tensor’s values, and a generic axis’s size used as a number. A function whose body has any of them, or that has errors, simply has no curve; there is no diagnostic. There is no ** operator: write pow(x, 3) or square(x).

The other parameters take the numbers written at the call, by position or keyword, also in a pipe stage (h |> act, h |> scaled(k=3.0)); otherwise their defaults, and a sentence says when an argument isn’t a number written in the source. A host scalar without a default and without a number at the call (as where the function is declared) means no curve. The formula is written from the body, with those numbers in it; a body with an if, or whose formula would be too long, shows the body as written instead. The range is the interval interpreter’s for the call’s numbers, narrowed by the checker’s domain for what the body returns and, for an elementwise function, by its declared return type.

Pictures of what a function does, on hover

In the editor, hovering a function’s name, or a call of a built-in that computes (argmax(z, over=C), a |> into softmax, -> in einsum’s or reduce’s pattern, x to T between domains), shows a picture of what it does: the language server runs it on a small example and draws what happens. Nothing is drawn by hand for any function; the example and the drawing come from the types, so a function of your own gets a picture as a built-in does. A function of one number has its curve instead.

The example. The file’s literal dims are made small (dim D = 768 becomes 6, or 4, 8 or 12 if 6 doesn’t check), generic axes get small sizes (the first 3, the others 5, a group’s parts 2), and each tensor parameter gets values its type allows, from a seeded generator, so a picture is the same each time: logits and reals from a normal distribution, a Prob over m a softmax along m, an Index[C] whole numbers below C, a Mask true or false, Unit from 0 to 1, u8 from 0 to 255. A model parameter is made with init(seed=0), a Mode is Eval, and a number without a default is 1 (int 2, bool true). A call’s picture runs a function that takes each of the call’s arguments that depends on the caller’s values, with the type the checker gave it, and makes the call with the rest as written, so masked_fill(!causal[n m], -inf) keeps its mask.

What’s drawn. One slice of each tensor, index 0 along all but its last two axes (at b = 0), up to 8 rows and 10 columns. Parameters whose sizes are all fixed, such as w: [D C], are weights, and aren’t drawn; the caption names them. Then:

• When the result’s type says something about its values (a Prob or LogProb over an axis, an Index, a Mask, a structure such as a mask’s triangle, or one number), when an input is an index, when the result has −∞ where a mask was, or for a conversion with to, the values going in and coming out are drawn, brighter for larger, from each grid’s smallest to its largest, which the grid says (−1.58 … 1.93). A Prob row is shaded from 0 to 1, with its sum, Σ 1; an Index of one axis is a 1 where it points, among 0s; −∞ is a dashed cell. An index input that picks values of another input, as cross_entropy’s y does of z, is drawn as those values, outlined.
• Otherwise, as for anything made of matrix products, whose values on a random example are noise, what depends on what is drawn: one value of the first data input, near the start, is changed a lot, up and down, and the values of the result that change are outlined. A value that doesn’t depend on it comes out bit for bit the same. Causal attention changes at that position and after it; an MLP only in its row; a normalization along its whole row; a running total from there on.

When there’s no picture. The function is run only if the client allows it (the pictures initialization option, or the shapescript/setPictures notification with {enabled}; VS Code allows it in a trusted workspace, unless shapescript.hover.pictures is off), only in the Python interpreter the server was given (the python option or shapescript/setPython), which needs PyTorch, and only for a function that, with every function of the file it calls, doesn’t loop over a dataset, print, save, load, write logs, use Python or call an extern fn that isn’t PyTorch’s or the prelude’s, nor takes a mut parameter, a dataset, logs or optimizer state. A file that imports other .ss files, or has errors, isn’t run. A hover waits up to 4 seconds; an answer that comes later is kept for the next hover.

The interpreter runs a worker the first time it’s needed, which imports PyTorch once and runs each function as standalone Python compiled from the file with its small dims, on the CPU, with gradients off. VS Code draws the picture as an image in the theme’s colors; other editors get it as text, with ·, ░, ▒, ▓ and █ for the values and ● for what depends on what.

A function’s formulas in index notation, which the language server can also write from its body (crates/lsp/src/formula.rs), aren’t shown on hover.

Broadcasting by name

The order rule

Elementwise operations (+ - * /, comparisons, and, or, and elementwise built-ins and functions) line their operands’ axes up as follows. A rank-0 operand broadcasts to anything. Otherwise:

1. By name. If every axis of one operand has a name and those names are all axes of the other operand, the axes line up by name. The result has the other operand’s axes, in its order. If both operands qualify, as with [n m] + [m n], the left operand’s order wins. The matched axes must have equal sizes (E0401). If a name appears several times, the ith axis with that name in one matches the ith in the other.
2. From the right. Otherwise the axes line up from the right, as in NumPy. Each pair must have the same size, or one of them must have size 1, which broadcasts (E0401). Two axes in a pair that both have names must have the same name (E0404). The result takes the size of each pair (a named axis is kept over an unnamed one of the same size), plus the extra leading axes of the longer operand.

So [b H n m] + [n m] lines up by name and gives [b H n m], as does [n m] + [b H n m]. [n D] + [k D] is E0404 even when n and k have the same size: rename one with rearrange or as if they really are the same axis. [n D] + [D] lines up by name; [n 3] + [3] lines up from the right.

The generated Python permutes and unsqueezes operands so they line up as the checker says, so no unsqueeze appears in the source. masked_fill requires its mask to broadcast to the tensor’s shape without adding axes (E0404).

Matrix multiplication

x @ y and matmul(x, y) work like torch.matmul. The last axis of x is contracted with the second-to-last axis of y, or its only axis if it has one. The contracted axes need equal sizes but not equal names (E0401), so [b n (H Dh)] @ [D D] is fine when H*Dh = D. With two or more axes each, the axes before the last two broadcast by the order rule, and the result is those, then x’s second-to-last axis, then y’s last. A rank-1 operand loses its contracted axis, and a rank-0 operand is E0403.

Both operands must have the same dtype (E0502), since torch doesn’t promote for @; the result is Real.

Operators

An operand that is a Python value makes the whole result a Python value (Python interop). A string operand is E0412.

Pipes

a |> f(args) means f(a, args), and a |> f means f(a). In a stage that is a call, a _ argument (positional or keyword) marks where the input goes instead: x |> add(y, _) means add(y, x). A stage may have at most one _ (E0205).

a |> to T means a to T (Conversions with to).

|> binds more loosely than every other operator. A line that starts with |> continues the line above (Lines and indentation):

y = x
    |> layer_norm(over=D)
    |> matmul(w1)
    |> gelu

Each stage is checked with the previous stage’s result as its input. A diagnostic inside a stage gets a note “in stage N of this pipeline, whose input is [...]”, and the editor shows each stage’s shape.

The stages map, filter, zip and enumerate iterate their input, which can be a size or a range as well as a tensor: see Iterator chains. In a pipe stage, these four names are always the iterator stages.

In the editor, each pipe line ends with an inlay hint of its shape, and hovering the hint shows the type with its diagram. Hovering |> into a built-in that only moves values, such as rearrange, or into a function of your own shows the value before the stage and after it, and hovering -> in a pattern of such a built-in shows the operands and the result: side by side, with an arrow between them. Into a built-in that computes, such as softmax, or on -> in einsum’s or reduce’s pattern, it shows the types before and after on a line and a picture of what it does instead. VS Code draws them as one image; other editors get the text drawings next to each other, with → between them, or one under the other if that would be wider than 120 columns.

Pattern operations

rearrange, reduce, repeat, einsum and split take a pattern inputs -> output. Each input side describes one operand: it must have as many entries as the operand has axes (E0403), and its entries bind by position:

• A name already in scope (a dim, a parameter’s axis, an earlier pattern’s name) must have the axis’s size (E0401).
• A new name is introduced with the axis’s size, and stays in scope for the rest of the function once the operation checks.
• A number must equal the axis’s size (E0401).
• A group (gh p1), or a grouped size 16gh, which is (gh 16), binds one axis whose size is the product of the parts. At most one part may be unknown, and it is solved for (E0402, E0409); a group that can’t fit is E0401. Groups can’t be nested, and a name can appear only once per side (E0407).
• A name that is the name of a different axis of the operand, as D in rearrange(x, b D n -> ...) on a [b n D], is E0408: patterns name axes by position, and this is probably a transposition.

Keyword arguments give sizes to new names: rearrange(x, b n (h d) -> b h n d, h=12). The value is a positive integer, or an integer value known at run time, which gives an unknown size. Naming a size that’s already in scope, or a name that isn’t in the pattern, is E0406.

The output side’s names must be names from the inputs (E0405), except in repeat. A number on the output other than 1 must match an equal number on the input, except in repeat (E0407). The result’s axes are named after the output’s names; groups and numbers give unnamed axes.

A pattern that renames the axis a Logit, Prob or LogProb domain is over carries the domain to the new name; grouping that axis leaves just its range. reduce gives the element type its reduction would (Built-in operations).

Indexing

x[i1, i2, ...] applies one index per leading axis of x; axes without an index are kept. More indices than axes is E0403.

Anything else as an index is E0406.

Slices. Each part is optional and must be an integer (E0406). When the bounds and the step are sizes (numbers, dims, axis names, and sums and products of them), the slice has an exact size if the checker can prove that it fits in the axis, assuming that axes aren’t empty (generic sizes at least 1, unknown sizes at least 0). A bound that is surely negative counts from the end, as in Python. So on an axis n, 0:n-1, 1: and :-1 all have size n-1, and 0:2*h:2 on (2 h) has size h. When every size is known, bounds past the ends are clamped as in Python: x[0:1000] on a [100 D] is [100 D]. Otherwise, as for 0:3 on an axis n that may be shorter than 3, or for bounds that are run-time values, the slice has a new unknown size, and no run-time check is added.

causal[n m] is a [n m] Mask that is true where the m index is at most the n index. Its two entries must be axis names in scope (E0301); anything else is E0406. Only causal takes axis names separated by spaces.

Built-in operations

Built-in operations that need a rule of their own are built into the checker; the rest are prelude stubs. Operations that work along an axis take it by name with over=. A missing over= where one is needed, or one that isn’t a name, is E0406, and a name that isn’t an axis of the tensor is E0405. An operation that needs an axis of a Python value is E0413.

In cat, stack, gather and scatter, the axes other than the over= axis line up by position, as in torch: they must have the same sizes (E0401), the same count (E0403), and two of them with names must have the same name (E0404). A Mask can’t be joined with numbers (E0501). topk and sort have no largest= or descending= keyword yet; negate instead.

Wrong numbers of arguments and unknown keywords are E0406. Arithmetic on a Mask is E0505, and a Mask where numbers are expected is E0501.

In the editor, hovering a built-in’s name in a call or a pipe stage shows its signature with named axes, what it does and an example, then its types before and after and a picture of what this call does, or, for a built-in that only moves values, this call’s tensors before and the result after, side by side. Go to definition on it opens its declaration: the prelude’s stub, or, for the operations in the table above that the checker types itself, builtins.ss, a file of declarations the language server generates from the checker’s list of built-ins and the documentation its hovers show. Its signatures are written as fn softmax(x: [... n] Logit over n, over: Axis[n]) -> [... n] Prob over n = builtin, where [... n] is any axes and the axis over names, which can be any of them, and types ShapeScript can’t write are names that say what they are (Function, Model, Pattern, Reduction, Type, Any). The file is only for reading: nothing reads it.

Assertions with as

x as T claims that x already has type T. T may give a shape, a domain, a dtype, or several. It is the only way to tell the checker something it can’t prove, and it always compiles to a run-time check (Runtime checks).

Statically, an as must be possible:

• A shape must have the same number of axes, and no axis may have a known size that differs from a known size of x’s (E0410). Symbolic sizes are accepted: rows as [196 D] on a [?k D] is fine. The result has the asserted axes, with their names.
• A domain must not contradict what is known about the values (E0410): a Mask asserted to be numbers or the other way round, a range that doesn’t meet the known range, Additive for values bounded below, Binary for values that can’t be 0 or 1, Index[V] for non-integers or a different V, or Prob for Logit over the same axis. Asserting a range keeps whatever is known that’s narrower: relu(x) as Unit is [0, 1] and sigmoid(x) as Pos keeps the sigmoid’s range where it’s narrower.
• A dtype must be of the same kind (bool, integer or float) as x’s (E0410); to change it, use to. A different dtype of the same kind is accepted.
• as [n] bool asserts a Mask, as x: [n] bool declares one.

The parts of the element type that T doesn’t write are kept: x as [b C] keeps x’s domain, and x as Unit keeps its shape. A domain written without a shape, such as x as Prob over C, finds C among x’s axes. Sizes in T must be in scope or written as ?k (Axis names).

A Python value, whose shape is unknown, gets the asserted type: torch.load("w.pt") as [D 3*D].

Value domains

The element type of a tensor has two parts: a dtype, how values are stored, and a domain, which values are allowed. Both are optional in a type. The domain defaults to Real, and the dtype defaults to the one the domain implies: bool for Mask, i64 for Index, and the module default float, f32, otherwise.

Built-in domains

over C must name one of the type’s axes, and only Logit, Prob and LogProb take it (E0503). Only Index takes a size in brackets. A combination of domain and dtype that makes no sense, such as Mask f32, Index[V] f32, Prob over C i64 or Unit bool, is E0503.

Values are assumed to be finite: NaN isn’t tracked, and an infinite end of a range is open, so -inf from masking leaves Real values Real.

Messages show a range by name when it has one (NonNeg, Pos, Unit, Signed, or a user domain), not at all for Real, and otherwise as an interval: [b C] [0, 2]. A dtype is shown only when it isn’t the one the domain implies: [b 3 h w] u8.

Widening

A value fits where a domain is expected when every value it may hold is allowed there. There is no conversion and no run-time cost.

• A range fits a range that contains it: Unit fits NonNeg, Signed, Real and any user domain that contains $[0, 1]$.
• Any numbers fit Logit over C: plain Real is implicitly a logit, because a linear layer’s output really is one. LogProb over C fits Logit over C too. Prob over C doesn’t.
• A semantic domain is matched by position: Prob over C fits Prob over K when C and K are the same axis position in the two types. Along any other axis it is just its range, so it fits Unit or NonNeg.
• Mask, Binary and Additive fit only themselves, and Index[V] fits only Index[V] with an equal V (E0507).

Domains are checked at calls (shapeless parameters too), at returns, and in the built-ins that need one. A value that doesn’t fit is E0501; its help suggests to T when the to table has a conversion, and as T otherwise.

Range inference

The checker tracks a closed or open interval for every numeric value, through elementwise arithmetic, activations, clamp and comparisons, so most domains are inferred: sigmoid(x) is Unit and sigmoid(x) * 2 - 1 is Signed. The rules:

• + - * / by the corners of the operands’ ranges; neg, abs, relu, clamp, maximum and minimum exactly.
• exp, expm1, log, log1p, sqrt, rsqrt, sigmoid, tanh, softplus and log_sigmoid at the ends of the range, with exact anchors such as exp(0) = 1; erf, gelu and silu by sign and bounds; sin and cos give $[-1, 1]$; floor, ceil and round give whole ends; pow only for exponents 1 and 2 or a non-negative base.
• sum, prod, cumsum and cumprod keep only the sign; max and min keep the range; mean keeps ends of -1, 0 and 1 and widens others slightly; std, var and norm give NonNeg; logsumexp is at least the lower end.
• @, einsum, layer_norm, convolution and anything else not listed give Real.
• A partial function (log, log1p, sqrt, rsqrt, the divisor of /) of values not proven to be inside its domain gives Real.

Ends are rounded outward to f32, so the ranges hold for the values torch computes. Ranges are approximate: x - x on Signed values is $[-2, 2]$, not 0. When the checker can’t prove a range, use as.

A call to an elementwise function gets its range from the function’s body, and a call to an extern fn whose Python function’s math the checker knows gets its range from that math (The prelude).

Semantic domains

softmax(z, over=C) gives Prob over C and log_softmax gives LogProb over C. Both take logits, so softmax of a Prob is E0501 (“softmax applied twice?”), and passing a Prob where a Logit is wanted says “softmax applied before cross_entropy?”.

Logits shifted or scaled by plain numbers stay logits, and so do negated logits, where, masked_fill, and reductions over another axis. Prob and LogProb become their ranges under arithmetic. mean over another axis keeps Prob over C, since a mixture of distributions is one. Removing the axis C (reducing it, indexing it away, grouping it in a pattern) leaves just the range.

Discrete domains

Mask comes from comparisons, !, and, or, any, all and causal. Where a mask is needed (!, and, or, the condition of where, the mask of masked_fill), anything else is E0501. Arithmetic on a mask is E0505; its help suggests to Additive for +, to Binary for the other operators, ! for -m, and and or or for two masks.

Binary times Binary, their maximum and minimum, and abs, relu and rounding of one stay Binary; other arithmetic gives $[0, 1]$. Additive plus Additive stays Additive.

Index[V] connects values to shapes: emb[ids] with emb: [V D] and ids: [b t] Index[V] is checked, so token ids from a different vocabulary are E0507. argmax, argmin, argsort and topk give indices. Plain integer indices are still allowed, unchecked.

Function edges

For log, log1p, rsqrt and the divisor of / and div:

• a range that only touches the edge, such as log of $[0, \cdot]$, or a divisor whose range contains 0 and isn’t all of Real, is the warning W0501;
• a range with a finite lower end that reaches past the edge (log of Signed), a range wholly outside, or a divisor that is exactly 0 is the error E0506.

sqrt only has the error, since sqrt(0) is fine. A range unbounded on that side, as for Real, gives nothing, because nothing is known. log of a sigmoid suggests log_sigmoid of its input, and log of a Prob suggests log_softmax.

User domains

domain Temp = [0.05, 10] declares a range domain. [ and ] include an end, ( and ) exclude it. The ends must be numbers, possibly negative, or inf (E0503), the range can’t be empty, and the name can’t be a built-in domain or a dtype (E0503). Each end covers the number as written and its nearest f32. A user domain works wherever a built-in range does, and is shown by name.

Conversions with to

There are three ways to change a domain, and each has one meaning:

The to table

A to with a domain makes one of these conversions, chosen from x’s type:

Anything else is E0504, which lists the table. That includes semantic and discrete domains matched against range rows by their values (Prob to Signed; assert first, p as Unit |> to Signed), and converting to the domain a value already has (“nothing to convert”).

u8 to Unit and Mask to Binary or Additive give the default float dtype; the others keep the float dtype. A to may also write a dtype, as in to Unit f16, which must be a float (E0504). A to with a domain on a Python value is E0413, because the conversion depends on a type the value doesn’t have yet. The shape in to [b C] Unit, if given, must be x’s (E0401).

Converting dtypes

A to with only a dtype changes how values are stored and keeps them:

• bool to a float gives Binary, and to an integer gives $[0, 1]$;
• an integer to a float keeps the range; a float to an integer truncates it;
• a float to a lower-precision float rounds the range’s ends outward;
• anything but bool to bool is E0504: compare instead, as in x != 0.

A dtype-only to on a Python value is allowed and leaves it a Python value.

Dtypes

The dtypes are bool, u8, i8, i16, i32, i64, f16, bf16, f32 and f64. Numbers in the source, true and false, and sizes used as values are host scalars: Python int, float and bool. The module default float dtype is f32; there is no syntax to change it yet.

Promotion follows torch’s result_type: bool < integer < float; tensors with axes outrank zero-dimensional tensors, which outrank Python numbers, unless a lower tier is of a higher kind; a Python float counts as the default float. So [b D] bf16 * 0.5 stays bf16, [b D] u8 * 2 stays u8, and bf16 + f16 is f32. / and float functions of integers give the default float, and sum and prod of booleans and integers give i64.

@, matmul and einsum need one dtype, since torch doesn’t promote them (E0502). Calls and returns compare dtypes exactly (E0502), except that a Python number fits any dtype of its kind or a higher one, a shapeless parameter takes any dtype unless it names one, and the tensor parameters of an extern fn take the call’s float dtype (Extern declarations).

Extern functions and the prelude

Extern declarations

## The complementary error function.
extern fn erfc(x: Real) -> NonNeg = "torch.special.erfc"

An extern fn gives a Python function a typed signature. The string is the function’s dotted Python name: a module and a function, each part a Python identifier and not a keyword (E0308). An extern fn has no body, so it needs a return type (E0308).

Calls are checked against the signature like calls to any function, and compile to calls to the Python function, with its module imported. The signature is trusted: what the function returns isn’t checked at run time.

A tensor parameter whose type doesn’t name a dtype takes any float dtype, the same for all such parameters in a call, and the result has it too, so a bf16 model stays bf16 (E0502 for a mix). A tensor parameter whose default is none may be left out, and the Python function then gets None.

An extern fn’s signature can say more than a fn’s, because there is no body to check. The next three sections describe what.

Rest axes

... in a type stands for any number of other axes, the same ones in every type of the signature. The result’s ... is the axes the argument’s ... had, names included:

extern fn roll(input: [...], shifts: Size[S], dims: Axis[n]) -> [...] = "torch.roll"
extern fn linear(input: [... Din], weight: [Dout Din]) -> [... Dout] = "torch.nn.functional.linear"

A type can have one ..., and it can’t be inside a group (E0306). Only an extern fn can use it (E0306). A return type with ... needs a parameter with one (E0308).

Axis and Size parameters

• dims: Axis[n] takes the name of an axis of the first argument, as in roll(x, 2, dims=t), whose size is n. The Python function gets the axis’s position: torch.roll(x, 2, dims=1). As anywhere in a signature, a name that is a dim is that size, not a new name: with dim N = 16, Axis[N] takes only an axis of size 16, and any other is E0401, whose help says so. The first parameter must have a shape, and an Axis parameter can’t have a default (E0308). A name that isn’t an axis of the first argument is E0405; an argument that isn’t a name is E0406. This also gets around dim being a keyword: the parameter can have another name, and be passed by position.
• dims: Axes[n] is the same, except that the Python function gets a list of the position, dims=[1], for the functions that only take a sequence of axes, such as torch.flip: extern fn flip(input: [...], dims: Axes[n]) -> [...] = "torch.flip".
• shifts: Size[S] takes a whole number known before the program runs (a number, a dim or an axis name), which is the size S (E0406 otherwise). A default, which must be a number or a size from an earlier parameter (E0308), is used when the argument is left out: stride: Size[S] = K.

All three are E0306 in a fn.

Rounded division

A division in an extern fn’s return type that isn’t exact rounds down, as in torch’s formulas for convolution and pooling:

extern fn conv2d(input: [b Cin H W], weight: [Cout Cin kh kw], bias: [Cout] = none, stride: Size[S] = 1, padding: Size[P] = 0) -> [b Cout (H + 2*P - kh)/S + 1 (W + 2*P - kw)/S + 1] = "torch.nn.functional.conv2d"

Each such division is worked out at each call. Once the arguments determine every size, it is a number. With a size that isn’t known, such as a generic h, the division must be exact (E0402). A divisor of 0 is E0406, and a result with an axis smaller than 1 is E0401.

The prelude

The prelude is two files of extern fn stubs that every file sees, crates/check/src/prelude.ss and crates/check/src/nn.ss. A declaration of the same name in the file hides a stub.

• prelude.ss: exp, log, log1p, expm1, sqrt, rsqrt, sin, cos, tanh, sigmoid, erf, gelu, silu, softplus and log_sigmoid, each (x: Real) with a return type that bounds its values.
• nn.ss: linear, layer_norm, rms_norm, batch_norm, conv2d, max_pool2d, avg_pool2d, embedding, scaled_dot_product_attention and dropout.

Go to definition on a stub’s name in the editor opens the stub. The language server writes prelude.ss, nn.ss and the generated builtins.ss (Built-in operations) into a folder, read-only, and gives their file:// locations, which any editor can open, and it reports no diagnostics for them. The folder is the libraryDir initialization option if the client sends one (the VS Code extension sends its global storage, and shows the files as read-only documents of a shapescript-library: scheme), and otherwise shapescript/lsp-<version> in the user’s cache folder: ~/Library/Caches on macOS, %LOCALAPPDATA% on Windows, $XDG_CACHE_HOME or ~/.cache elsewhere. Neovim opens them read-only because they aren’t writable.

The checker knows the math of these Python functions: torch.exp, torch.log, torch.log1p, torch.expm1, torch.sqrt, torch.rsqrt, torch.sin, torch.cos, torch.tanh, torch.sigmoid, torch.erf, and torch.nn.functional‘s gelu, silu, softplus and logsigmoid. A call to an extern fn, in the prelude or not, that names one of them gets its range from that math and the arguments’ ranges, as precise as a built-in’s, and its edges are checked. Other stubs’ results have their return type’s range.

Host scalars

int, float and bool are host scalars: Python numbers and booleans, kept on the host rather than on the device, for counters, epochs, seeds and flags. A number written in the source is an int or a float, true and false are bools, and a dim or axis name used as a value is an int. A type that is only int, float or bool declares one, as in steps: int or -> bool. Internally a host scalar is a value of shape [] with a host dtype, and messages show it as int, float or bool.

• Arithmetic on host scalars gives a host scalar: int with int gives an int, except that / gives a float, and anything with a float gives a float. Arithmetic with a tensor gives a tensor, with the dtype promotion gives.
• A comparison of host scalars gives a bool, and and, or and ! take and give bools; they compile to Python’s and, or and not. A bool isn’t a number, so arithmetic on one is E0505, and !, and or or on an int is E0501.
• A parameter or return type that is a host scalar takes a host scalar or a Python value. An int fits where a float is expected, and nothing else converts: a float where an int is expected is E0502, and a tensor, even one of shape [], is E0805, since reading it would wait for the device.
• x as int, x as float and x as bool give a Python value a host type, with a run-time check that it is one (Runtime checks); a bool is neither an int nor a float. On a tensor, such an as is E0805. to int is E0504: no conversion makes a host scalar of a tensor.
• An int or float has a range, as a tensor does (Range inference), so a counter that starts at 0 and is only increased in a loop is in $[0, \infty)$ (Blocks and scope).

A host scalar parameter isn’t checked at run time.

Control flow

Loops

for i in a..b: runs its block once for each int from a up to, but not including, b, as Python’s range(a, b) does. The ends must be ints or Python values (E0803): a range counts on the host, so its ends can’t be tensors, even [] i64 ones. The loop variable is an int whose range goes from a’s lowest value to b’s highest value minus one.

for item in e: loops over a Python value e, and each item is a Python value (Python interop), or over a dataset’s batches, as in for batch in train |> batches(64):, and each item is a batch (Looping over batches). Looping over anything else, such as a tensor or a dataset that isn’t in batches, is E0804.

A return in a loop returns from the function. There is no while, break or continue.

Conditions

if c: runs its block when c is true, and the else: block, if there is one, when it isn’t. c must be a bool, a [] Mask or a Python value (E0802). Reading a [] Mask, as in if sum(x) > 0: or if any(x > 0):, waits for the device to compute it. A mask with axes isn’t true or false; reduce it with any or all. A Python value is true or false as Python says.

Blocks and scope

The block of a for, and each block of an if, are scopes of their own:

• The loop variable, and the values and axis names first bound in a block, are out of scope after it (E0301).
• A new value that both blocks of an if with an else bind, with the same shape, axis names and dtype, is in scope after the if.
• A block may assign a value from outside it, which then has the new value after the block. The value must keep its shape, axis names and dtype, or stay the same model or a Python value (E0801), so that the code after an if and the next iteration of a loop work whichever way the block went. A host scalar keeps its kind too: adding 0.5 to an int counter in a loop is E0801. A loop whose variable has the name of a value from outside the loop assigns that value. To compute a value of another type, give it a new name.
• A value’s domain may change in a block. After an if, it is the narrowest domain that holds what either block leaves, such as the hull of two ranges. A loop’s body is checked repeatedly, without reporting, until the domains at its start hold those at its end. A range that still grows after two rounds has its growing ends moved to infinity, and a domain that still changes after six rounds becomes Real. So a counter that starts at 0 and only goes up is in $[0, \infty)$, and y = y * 2 in a loop leaves Unit values in $[0, \infty)$, which no longer fit a Unit return type (E0501).

Iterator chains

An iterator chain computes something for every element along an axis at once, as pipe stages: n |> map(i => e) is a tensor whose first axis is n, with e computed for every index i along it. It’s the parallel counterpart of a for loop, which stays sequential: a loop’s iterations run one after another and may depend on each other, while the elements of an iterator stage are computed together and can’t. No Python loop runs.

dim D = 8
dim C = 4

fn scale_rows(x: [n D], w: [n]) -> [n D] = n |> map(i => x[i] * w[i])

fn per_row(x: [n D], w: [D C]) -> [n C] = x |> map(row => softmax(row @ w, over=C), over=n)

fn gram(x: [n D]) -> [n n] = x |> map(a => x |> map(b => sum(a * b), over=n), over=n)

fn table() -> [3 D] f32 = 3 |> map(i => D |> map(j => j to f32))

fn positive(x: [n D]) -> [] = x
    |> map(row => sum(row), over=n)
    |> filter(t => t > 0.0)
    |> sum

What a chain iterates

The input of the first iterator stage says what is iterated:

• A size: an axis name of the function’s types, such as n in x: [n D], a loop’s batch axis b, a dim, a positive whole number (3), or a sum or product of them (2 * n). Its elements are its indices, each a [] i64 Index[n]: a tensor, not a host int, since every index is computed at once. So x[i] takes x’s first axis at i, as indexing with any Index[n] does (Indexing), and x[i] on an axis of another size is E0507; i to f32 is its value as a float. The iterated axis is named like the size when the size is a name, and has no name otherwise.
• A range a..b, whose ends are sizes: the indices from a up to b, an axis of size b - a. 0..n is the same as n; a range that doesn’t start at 0 gives [] i64 elements, not Index.
• A tensor: its slices along an axis, each with that axis removed. The axis is the one named by over= on the stage; or, if the tensor has one axis, that one; or, if the tensor is the result of a map or filter, the axis it iterated, which comes first in it. Otherwise the stage needs over= (E0406), and an over= that isn’t one of its axes is E0405.

Anything else, such as an int value known only when the program runs, a [] tensor or a Python value, is E2401: the result’s axis is the one iterated, so its size must be known from the types. A Python value can be given a shape with as first.

Stages

zip’s other iteration is a size, a range or a tensor, iterated along its first axis if it’s the result of a map or filter, its only axis, or the axis with the iterated axis’s name (E0406 otherwise). Zipping pairs with a third iteration gives triples. Pairs, from zip or enumerate, must be taken apart by a map or filter whose lambda has a parameter for each (E3206); anything else, such as ending a chain with them or giving them to another stage, is E3205. filter of pairs keeps pairs.

The stages compose like any pipe: a map’s result is a tensor that keeps the iterated axis’s name, so a reduction can name it, as in x |> map(row => relu(row), over=n) |> sum(over=n). After a filter, whose axis has no name, a reduction without over= reduces it, with the others. Each pipe line gets its inlay hint: a pair shows as its two tensors, each with the iterated axis first.

Functions

The function of map and filter is one of:

• A lambda, v => body or (a, b) => body, with a parameter for each part of an element (one, or two for pairs). Each parameter is in scope only in the body, can’t have the name of an axis or a dim (E0303), and _ ignores it. The body is an expression, which may be another chain: 3 |> map(i => D |> map(j => j to f32)) is [3 D] f32, with i and j both in scope in the inner body. A lambda anywhere else is E3202: lambdas aren’t values.
• A pipe stage, applied to each element as |> would apply it (Pipes): a function’s name (map(gelu), map(normalize, over=n)), a call that the element is passed to first or as its _ (map(clamp(0.0, 1.0)), map(sub(1.0, _))), or to T (map(to f32)). These take one element, not pairs (E3206).

Independence

No element may depend on another, and computing one may not change anything; the checker enforces both:

• A lambda can’t read the result it builds. In y = n |> map(i => y[i - 1] + x[i]), where y isn’t bound before, y would be other elements of the result: E2402. Each element computed from the one before is a recurrence, which a for loop computes. If y is bound before the statement, the lambda reads that old value, which is fine.
• The function only computes a value (E2403). It can’t use mut arguments, Python (a call through import python, or an extern fn outside torch and math), print, save or load, a mode (dropout and other random draws, whose random numbers would be shared or differ by order), value_and_grad, optimizers, update, schedules, datasets, logs records or sinks, or a slice bound that depends on a parameter, since every element must have the same shape. Reading models, their fields, gradients and batches is fine. A map must give a tensor, not a model.
• Functions it calls, in the lambda or given by name, are checked the same way, at any depth, and also can’t change a model’s field, branch with an if on a tensor (which would read the device for each element) or loop over a Python value. The error is at the call, with the construct in the function as a secondary label. A function with mut parameters is only called with mut arguments, which are reported.

The input of a chain, zip’s other iteration and over= are computed once, before the stage, and aren’t checked this way.

How iterator chains run

Each map and filter is one batched computation, chosen by the checker:

• Plain tensor operations, when a lambda is elementwise in its parameters: made of the parameters, x[i] (when i is the index of a size), numbers, values that don’t use the parameters, arithmetic and comparisons (not @), to conversions of the dtype or domain, and elementwise functions (built-in, extern fn or user functions whose parameters have no shapes), with every part that uses a parameter shaped like the body or []. The body is computed once on whole tensors: a parameter becomes the whole iterated tensor, x[i] becomes x, and a [] part gets size-1 axes to line up. So n |> map(i => x[i] * w[i]) is x * w[:, None]. A function given by name that is elementwise, or to T without a shape, is applied to the whole tensor: x |> map(gelu, over=n) is gelu(x).
• torch.vmap otherwise: torch.vmap(lambda row: body)(x), over the iterated tensors with the iterated axis moved first, nested for a nested chain.

A size’s indices are torch.arange(n), and a range’s torch.arange(a, b). filter computes its mask the same way, then indexes with it. zip and enumerate only pair tensors up: enumerate’s indices are torch.arange of the axis’s length.

Either way autograd follows it, so gradients flow through a chain, value_and_grad included. Shape checks in a lambda (as [D], and the argument checks of functions it calls) check each element’s shape. A value check in a lambda checks every element’s values together, and fails as any value check does (Runtime checks).

Export. torch.export records torch.vmap, so an export fn can use iterator chains, and the axis they iterate can stay dynamic (Export). A filter gives an unknown size, which export handles as it does unknown sizes. In an exported program, a value check inside a lambda isn’t kept: torch.vmap can’t give every element’s values to one assertion. Check the result after the chain instead.

Padding. A map over a padded axis computes each position on its own, so it needs nothing when its function uses values with the padded axis only as the elements it’s given, or as x[i] with i an index of the axis. Any other use of the padded axis in a stage’s function, such as summing over it inside, and a filter, zip or enumerate over it, keeps the axis from being padded (H2001).

Several devices. In a function that runs on the workers, an iterator stage that uses a value split across the devices, along the batch axis or along H with split H, is E2404: it has no lowering that works on each device’s part. One that uses only whole values, such as a model’s params, is computed whole on every device. Rewrite one over the batch with operations on whole tensors, which the checker splits.

Loops that could be iterator chains

A sequential loop for i in 0..n:, over a whole axis n, gets hint H2401 when its iterations don’t depend on each other: every statement of its body binds a value of that iteration, or adds a term to an accumulator from before the loop (acc = acc + e or acc = e + acc); no term reads an accumulator; i is used only as an index, x[i]; and nothing in it does more than compute a value, as for a stage’s function. The help suggests acc = acc + (n |> map(i => e) |> sum(over=n)). A loop whose iterations read what the last one computed, that uses i as a number, or that prints, gets no hint.

Models

Model declarations

model Name: followed by an indented block of fields declares a model, a record of the tensors that make up a program’s state:

## A two-layer classifier.
model Mlp:
    param w1: [D Hd] = normal(std=0.02)
    param b1: [Hd] = zeros
    param w2: [Hd C] = uniform(-0.1, 0.1)
    ## Subtracted from every input. Saved with the model, never trained.
    buffer mean: [D] = zeros

A param field is trained. A buffer field is saved with the model but not trained. A model and each of its fields can have a doc comment. A model’s fields have different names (E0302). Models, like functions, can be used anywhere in the file.

Fields

A field’s type has a shape made only of numbers and dims, so every model of a type has the same shapes: a generic name is E0301, and ?k is E0902. For a single number, write []. The other rules are E0902:

• A param holds floats, of any float dtype (f32 unless the type names another), and has no domain, since training can change it to any value.
• A buffer can have any dtype, and a range domain, Index[V] or Mask, but not a semantic domain, Binary or Additive.
• A field can’t be a host scalar, a type without a shape, or a model: models don’t nest. Nor can it be optimizer state or a Schedule.
• A field can’t be named init, a Python keyword, a name that starts with _, or an attribute of torch.nn.Module, such as train, eval, to or half.

Initializers

Every field has an initializer after its =, which gives it its first value:

Arguments are given by position or by name, as in a call. A missing initializer, one that isn’t in the table, the wrong arguments, a negative std and a lo more than hi are E0903. The arguments are Python numbers, and may use dims, as in normal(std=1 / D): floats, except that full takes an int for an integer field and a bool for a bool field (E0502). normal and uniform only fill float fields, and the values an initializer gives must fit the field’s domain (E0903): zeros can’t start a Pos buffer, nor uniform(0, 2) a Unit one.

Using models

A model’s name is a type, as in m: Mlp for a parameter or -> Mlp for a return type. Only a model of that name fits it; another model or a tensor is E0904. A Python value fits, and is checked when the function runs (Runtime checks). A model’s name in any other type, as in [D] Mlp, is E0303.

m.w1 reads a field, and has the field’s type. A field the model doesn’t declare is E0901. A model isn’t a tensor: arithmetic on one, or passing it for a tensor parameter, is E0412.

Mlp.init(seed=s) makes a model. seed is an int, given by name (E0406; E0502 or E0805 for something other than an int). init is all a model’s name can be used for as a value: Mlp alone and Mlp(...) are E0303, and Mlp.other(...) is E0301.

Compiled models

A model compiles to a Python class, a subclass of torch.nn.Module. Mlp.init(seed=0), called from ShapeScript or from Python, returns one, with each param registered as an nn.Parameter and each buffer with register_buffer. So its state_dict() holds exactly the declared fields, and it works with existing PyTorch code. Storing any other tensor or module on it raises AttributeError.

Seeds. There is no global random state. Each field is drawn from a CPU torch.Generator of its own, seeded with the first 8 bytes of the SHA-256 of "{seed}/{field}", where field is the field’s name. So the same seed gives the same values, and adding, removing or reordering fields doesn’t change the others. Random values are drawn as f32 (f64 for an f64 field) and then converted to the field’s dtype, so they are the same on any device. Every field is made on the declared device.

A model’s architecture in the editor

In VS Code, a model’s architecture is drawn in a tab of its own, Architecture: Mlp, beside the editor. Clicking a model’s name, where it’s declared, used as a type (m: Mlp) or in Mlp.init, or the name of a function that takes a model first (fn forward(m: Mlp, …), or a call of one), shows that model there; the tab follows the cursor from model to model without taking the focus. It opens by itself the first time; once closed, it stays closed until Show architecture above a model line, the hover’s Open architecture link, Show Model Architecture in the editor’s context menu or the command palette opens it again (shapescript.architecture.followCursor: false stops it following the cursor). In the tab, the drawing is fitted to the width and can be zoomed (the buttons, Ctrl/Cmd + scroll, + - 0) and dragged; Layers and Every operation switch the level; the functions it was read from are links to their lines; and clicking a box selects the code it was drawn from. The hover then sums the architecture up in one line, from its layers, Architecture of Gpt: Embedding → Transformer block × 6 → LayerNorm → Linear · tied · Open architecture, above the model’s code and docs.

Other editors show the drawing on hover instead. Hovering a model’s name, where it’s declared or used as a type (m: Mlp), draws its architecture, read from the functions that take it as their first parameter, in its file and the files it imports (overloads): its forward, then its loss. The hover starts by naming them, with links to their lines: “Architecture of Mlp, from forward (mlp.ss:26) and loss (mlp.ss:29)”. Without a forward, it uses the other functions that take the model first and return a tensor (at most two, not those that take it mut); with none, it lists the fields, each with a small drawing of its shape.

The drawing goes from top to bottom, one box per layer, as papers and explainer videos draw a model: Embedding V → D, a Transformer block × L = 6 holding LayerNorm, Multi-head attention · H heads · causal, a residual connection ⊕, LayerNorm, MLP D → Dff → D · GELU and another ⊕, then LayerNorm, Linear D → V · tied to wte and the loss, Cross-entropy over V. Under each layer’s name is the type of its result, as the checker worked it out, with sizes named as the code writes them; the function’s inputs come first and its declared return type last. The params a layer holds are inside its box, each a small rectangle with the proportions of its diagram, its name and its type (just the names when they don’t fit); other inputs from the side, such as the labels cross_entropy compares with, are beside it, with an arrow in.

Layers are found from the operations, however the code writes them:

• Linear: a matrix product with a 2-D param, x @ m.w, matmul(m.w), linear(x, m.w, m.b) or an einsum of one input with a 2-D param, with a 1-D param added after it as its bias: Linear D → C. An activation right after it is part of it, Linear D → Dff · GELU; if its weights are an embedding’s, it says so, · tied to wte.
• MLP: linear layers with an activation after each but the last, and maybe dropout between them: MLP D → Dff → D · GELU, · dropout 0.1.
• Attention: scaled_dot_product_attention, or a softmax over a product of two computed values (einsum(q, k: …)), with the projections that make the queries, keys and values, the steps between (splitting heads, scaling, masking, the product with the values, merging heads) and the projection after: Multi-head attention · H heads, with · causal if a causal mask is applied. H is the axis before the last two of its scores or result.
• Embedding (embedding(ids, m.wte)), position embedding (a param of two or more axes added to every example), LayerNorm, RMSNorm, BatchNorm, GroupNorm (the built-ins, or a function named so that does arithmetic, as rms_norm written out does), convolutions (Conv k×k · C → D, or Patch embedding P×P · 3 → D when its stride is its kernel), dropout (a small tag, Dropout 0.1) and losses (Cross-entropy over V, MSE loss).
• A function that’s called is a frame named after it, unless it’s one layer, which is drawn as that layer; one that holds attention, a feed-forward layer and a residual connection is a Transformer block, with the function’s name on the right.
• Reshapes that only move axes (rearrange, split, transpose, casts) are left out, as diagrams usually leave them out, except a rearrange that cuts an image into patches, Patchify 16×16. Other operations stay boxes of their own, with readable names: Mean over n, Softmax over m, × 0.5.

A value that isn’t the line before’s, such as a residual connection’s x + a, is an arrow down the left, into a + circle. Calls that pass the model or its params to another function are followed, three calls deep, with a link to each in the header; a loop’s body is drawn once, as a frame marked × L (× L = 6 when L is a dim), and so are calls of one function in a row on each other’s results (× 3). Functions given only tensors, and forward called from loss, are one box each. Past 32 boxes, the innermost frames become one box each, and then the last boxes are counted instead of drawn. Layers are colored by kind, in the theme’s chart colors: blue for linear layers, MLPs and embeddings, green for activations, purple for norms, orange for attention and softmax, red for the loss.

The operations level draws every operation instead, as the code does it: matmul, add, layer_norm over D, split b n (3 H Dh) → b H n Dh, each with the params it reads to its right, and each function that’s called in its own frame. In VS Code, the setting shapescript.architecture.detail chooses the level (layers or operations), and the hover’s link, Show every operation or Show layers, switches it and shows the hover again. Other editors choose it with the initialization option architectureDetail: "operations", or later with the notification shapescript/setArchitectureDetail with {detail}; a client that runs command: links in hovers can ask for the link with architectureToggle: true, whose command, shapescript.architecture.toggleDetail, is the client’s to run.

A client can show the drawing in a view of its own, as VS Code’s tab does, with the request shapescript/architecture, given {textDocument: {uri}, position} and optionally detail ("layers" or "operations"). It answers null where there’s no model (a model’s name, or a function that takes a model first), and otherwise {model, title, summary, detail, svg, width, height, text, declaration, links, nodes}: summary is the one-line summary; svg is drawn in currentColor and the accents above, with each box in a <g data-node="id">; text is the list below; declaration is the model’s {uri, range}; links are the functions read (kind: "source") and followed ("followed"), each {label, place, kind, uri, range}; and nodes are {id, uri, range}, where each box’s code is. probe: true answers only {model}. A client that says so with architecturePanel: true gets the summary and a link to the command shapescript.architecture.show in the hover, rather than the drawing.

Editors that show images in hovers show the drawing as an image; a hover is short, so drag its bottom edge down to see all of a large model. Other editors get a list, one line per layer, with ↓ where it reads the line above, ⊕ for a sum with a value from further up, + x naming that value, the params a layer holds after its type, those beside it after ←, and frames drawn with ┌ │ └:

forward  mlp.ss:26
  • x                                    [b D]
  ↓ MLP D → Hd → C · GELU · dropout 0.1  [b C]
        w1 [D Hd], w2 [Hd C]
  ↓ output                               [b C] Logit over C

loss  mlp.ss:29
  • batch.img                            [b D]
  ↓ forward (above)                      [b C] Logit over C
  ↓ Cross-entropy over C                 []
        ← batch.label [b] Index[C]
  ↓ output                               [] NonNeg

and, at the operations level:

forward  mlp.ss:26
  • x                     [b D]
  ↓ matmul                [b Hd]              ← m.w1 [D Hd]
  ↓ gelu                  [b Hd]
  ↓ dropout 0.1           [b Hd]
  ↓ matmul                [b C]               ← m.w2 [Hd C]
  ↓ output                [b C] Logit over C

Only the first branch of an if is drawn. examples/gpt.ss is a GPT whose hover shows its embedding, its transformer block, with the attention and the MLP inside it, once, × L = 6, and its output layer, tied to the embedding.

Changing models with mut

mut parameters

A function may change a model, or optimizer state, only if it is a mut parameter, and every call writes mut before that argument, so both the signature and the call show what changes:

fn step(mut m: Mlp, x: [b D], y: [b] Index[C], lr: float):
    l, g = value_and_grad(loss, m, x, y)
    m.w1 = m.w1 - lr * g.w1
    m.w2 = m.w2 - lr * g.w2

fn fit(mut m: Mlp, x: [b D], y: [b] Index[C], steps: int):
    for i in 0..steps:
        step(mut m, x, y, 0.1)

Only a parameter whose type is a model or optimizer state, such as AdamW[Mlp], can be mut; a tensor, host scalar or schedule is a value, and a function returns a new one instead (E1101). An extern fn’s parameters can’t be mut (E1101). At a call, an argument for a mut parameter needs mut (E1103), and mut before an argument whose parameter isn’t mut, or before an argument of a built-in or of Mlp.init, is E1101. Only a name can be passed as mut (E1101), so the input of a pipe can’t be. A Python function can change anything, so an argument of one may be marked mut; what it passes must be something the caller may change.

What counts as a change

A model is changed by assigning one of its fields, m.w1 = e, or by passing it as mut, as to update; optimizer state only by passing it as mut. Reading fields, and passing a model to a parameter without mut, don’t change it. e must have the field’s shape and dtype (E0403, E0401, E0502), and fit a buffer’s domain (E0501); only a model’s fields can be assigned (E0303), and only declared ones (E0901).

A function may change:

• its mut parameters;
• models and optimizer state it makes or gets from calls, such as Mlp.init(seed=0) or adamw(m, lr=3e-4), unless they may be a parameter without mut (below);
• in a notebook cell, outside any fn, any value.

Anything else is E1102: a parameter without mut, and a value that may be one, because it was assigned from one (a = m) or computed from one with the same type (a function can return its argument). After an if or a loop, a value that may be such a parameter on some path is treated as one. A Python variable seen inside a function can’t be changed either.

A mut parameter can’t be given a new value, as in m = Mlp.init(seed=1) or by a for, since the caller would still have the model or state it passed (E1104).

Changes at run time

A field assignment compiles to rt.set_field(m, "w1", e), which copies the new value into the field’s tensor in place, under torch.no_grad(). So from Python, a mut argument is changed in place: the caller’s model object changes, its params stay the same nn.Parameter objects, and optimizers and state_dict() still refer to them. The new value isn’t part of any gradient computation.

Gradients

value_and_grad

value_and_grad(f, m, args…) calls f(m, args…) and gives two results, its value and its gradient with respect to the params of m, so it’s bound to two names (E0411):

fn loss(m: Mlp, x: [b D], y: [b] Index[C]) -> [] NonNeg:
    return cross_entropy(forward(m, x), y, over=C)

fn grads(m: Mlp, x: [b D], y: [b] Index[C]) -> Grad[Mlp]:
    l, g = value_and_grad(loss, m, x, y)
    return g

f is the name of a fn (not an extern fn, a built-in or a value) whose first parameter is a model, and which has no mut parameters (E1105). The other arguments are checked as for a call f(m, args…); m must be the model of f’s first parameter (E0904). f must return a [] of floats, such as a loss (E1106), and the value has that type. f may take a mode, as a loss with dropout does: value_and_grad(loss, m, x, y, Train(seed=step)); gradients go through the values dropout keeps.

Which params f reads is worked out from its code: m.w1 reads w1, and passing m to another fn reads what that one reads, also through pipes and nested value_and_grad. Passing m anywhere else, such as to Python, may read any param. If f reads no param of m, its gradient would always be zero, which is E1107; each param it never reads gets the warning W1101.

Gradient types

Grad[M] is the type of a gradient with respect to the model M. It has a field for each param of M, with the param’s type, read as g.w1. Buffers aren’t differentiated, so it has no fields for them (E0901). Grad[M] can be a parameter or return type; a Grad of something other than a model is E0303, and so is as Grad[M].

Gradients of the same model can be added and subtracted, -g negates one, and a gradient can be multiplied by a [] or a number on either side, or divided by one. norm(g) is the norm of all the gradient’s values together, a [] NonNeg. Anything else, such as g * g, g + x for a tensor x, 2 / g or mean(g), is E1108; adding gradients of different models is E0904.

Gradients at run time

value_and_grad compiles to rt.value_and_grad(f, m, args…), which runs f with autograd enabled, even inside torch.no_grad(), and calls torch.autograd.grad with the model’s params. Nothing reads or writes .grad, so there is nothing to zero, and no train() or eval() to switch; randomness comes from a mode argument, which value_and_grad passes on to f like any other. The value is detached. A gradient is an rt.Grad: its fields are attributes and items (g.w1, g["w1"]), in declaration order, and it supports the arithmetic above. Each gradient tensor has its param’s shape and dtype, and is where the param is, the declared device; a param that f doesn’t use gets zeros. A Grad[M] parameter checks its argument is a gradient of M with the params’ shapes and dtypes (Runtime checks).

Optimizers

Optimizer state

sgd(m, lr=…) and adamw(m, lr=…) make optimizer state for the model m. Its type comes from the model’s: adamw(m, lr=3e-4) with m: Mlp is an AdamW[Mlp], and sgd(m, lr=0.1) an SGD[Mlp]. It covers every param of the model, so none can be left out, and holds what the optimizer keeps between steps (AdamW’s moments, SGD’s momentum) and how many updates it has made.

fn fit(mut m: Mlp, x: [b D], y: [b] Index[C], steps: int) -> AdamW[Mlp]:
    opt = adamw(m, lr=warmup(100, cosine(3e-4, steps)))
    for i in 0..steps:
        l, g = value_and_grad(loss, m, x, y)
        update(mut m, mut opt, g |> clip_norm(1.0))
    return opt

The model comes first, by position; anything else there is E1501. The hyperparameters are Python numbers, given by name (E0406 for one given by position, an unknown name, or a missing lr), with PyTorch’s defaults:

lr is a learning rate, a number or a Schedule, and has no default. The other parameters are host scalars: a tensor is E0805 and another kind of number E0502, as for any float or bool parameter. Their values are checked by PyTorch when the state is made, so a negative lr raises a ValueError then.

AdamW[M] and SGD[M] are types, for parameters and return types. Optimizer state is matched by optimizer and model: state of another model’s type is E1502, and anything else, such as SGD[Mlp] where AdamW[Mlp] is expected, E1503. AdamW or SGD of something other than a model, as AdamW[M] and the like are E0303. A parameter whose type is optimizer state can’t have a default (E0306), and optimizer state can’t be a field of a model, a data record or a logs record. Like a model, it is changed only through a mut parameter or a value the function made (Changing models with mut); a mut parameter’s state can’t be replaced with new state (E1104).

update

update(mut m, mut opt, g) applies one step of the optimizer state opt to the model m, with the gradient g, at the learning rate opt’s schedule gives for the updates it has made so far, and counts the update. It changes m and opt, so both are written mut (E1103) and must be changeable here (E1102); g isn’t changed (mut g is E1101). It takes its three arguments by position (E0406) and has no result.

The model, the state and the gradient must be of the same model type: opt or g of another model is E1502, comparing with m’s type (or opt’s, when m is a Python value). A first argument that isn’t a model or a third that isn’t a gradient is E1505, and a second that isn’t optimizer state is E1503.

Gradient transforms

A gradient transform is a function from Grad[M] to Grad[M], used with pipes: g |> clip_norm(1.0). clip_norm(g, max_norm) scales the gradient so that the norm of all its values together, norm(g), is at most max_norm: by min(1, max_norm / (norm(g) + 1e-6)), as torch.nn.utils.clip_grad_norm_ does. max_norm is a float, by position or by name. Its first argument must be a gradient (E1505). A transform of one’s own is an ordinary fn, using the gradient arithmetic:

fn halve(g: Grad[Mlp]) -> Grad[Mlp] = g * 0.5

fn step(mut m: Mlp, mut opt: SGD[Mlp], g: Grad[Mlp]):
    update(mut m, mut opt, g |> clip_norm(1.0) |> halve)

Schedules

A learning rate is a number or a schedule, a value of type Schedule that gives one for each step. Schedules are made with:

• cosine(peak, steps, end=0.0): $\mathit{end} + (\mathit{peak} - \mathit{end}) \cdot \frac{1 + \cos(\pi \min(t, \mathit{steps}) / \mathit{steps})}{2}$ at step $t$. It starts at peak, follows a half cosine down to end at step steps, and stays at end after. Up to steps it’s torch.optim.lr_scheduler.CosineAnnealingLR with T_max=steps and eta_min=end, which would rise again after.
• warmup(steps, then): rises in a straight line from 0, $\mathit{then}(0) \cdot t / \mathit{steps}$ for $t < \mathit{steps}$, then is then started afresh, $\mathit{then}(t - \mathit{steps})$. then is a number or a schedule, so warmup(100, cosine(3e-4, 10000)) warms up for 100 steps and decays over 10000 more.

peak and end are floats and steps an int, by position or by name (E0406, E0502, E0805); steps must be at least 1 for cosine and 0 for warmup, which is checked when the schedule is made. then and lr= take a number or a schedule; anything else, such as a [] tensor, is E1504. Schedule is also a type, for parameters and return types, and a number is accepted where a Schedule is expected, as a constant learning rate. A schedule supports no arithmetic.

The step count lives in the optimizer state. opt starts at 0 updates; each update reads the schedule at the number of updates made before it, then adds one. So the first update uses the learning rate for step 0, as a PyTorch scheduler’s first optimizer.step() does, and a schedule is the same value however many optimizers use it. Saving the state saves the count and the schedule.

Optimizers at run time

sgd and adamw compile to rt.sgd and rt.adamw, which make an rt.Optimizer around torch.optim.SGD or torch.optim.AdamW over the model’s params, in field order. update compiles to rt.update(m, opt, g): it sets the learning rate from the schedule, gives each param its gradient as .grad, calls the PyTorch optimizer’s step(), and sets .grad back to None, even if the step fails, so no .grad is left behind and the params stay the same nn.Parameter objects, changed in place. So a step is exactly PyTorch’s. The state is per param of the model’s type, not per model object: given another model of the same type, update moves the state to that model’s params. clip_norm is rt.clip_norm, which makes a new rt.Grad.

AdamW is made with capturable=True where PyTorch supports it, so that a step can be recorded as a graph (FAST-1): PyTorch allows it on CUDA (and XPU, HPU and XLA devices), but not on the CPU or Apple’s MPS, where the state is made without it. A capturable optimizer keeps its learning rate in a device tensor, which update fills before each step. torch.optim.SGD has no capturable option and keeps no step count of its own, so it needs none.

opt.state_dict() gives everything needed to continue: the optimizer ("sgd" or "adamw"), the model’s type and param names, the learning rate (a number, or a schedule as a dict such as {"schedule": "cosine", "peak": 0.0003, "steps": 1000, "end": 0.0}), the update count, and PyTorch’s own state dict, with the learning rate as a number. opt.load_state_dict(state) restores it into state of the same optimizer and model type, and raises ShapeError for any other. opt.steps is the update count, opt.lr the learning rate and opt.torch the PyTorch optimizer. From Python, shapescript.rt.adamw(m, lr=…), shapescript.rt.sgd(m, lr=…), shapescript.rt.cosine(…) and shapescript.rt.warmup(…) make the same values. A parameter whose type is optimizer state checks that its argument is state of that optimizer for that model type, and a Schedule parameter that its argument is a schedule or a number (Runtime checks).

Modes and randomness

Modes

Mode is the type of a mode: Train(seed=s), for training, whose random operations draw from the seed s, or Eval, for evaluation, whose random operations do nothing. Random operations, of which dropout is the first, take a mode as an argument. There is no global training flag and no global random state, so a function that draws random numbers says so with a Mode parameter, and every call to it passes one:

fn forward(m: Mlp, x: [b D], mode: Mode) -> [b C]:
    return x @ m.w1 |> gelu |> dropout(0.1, mode) |> matmul(m.w2)

fn train_step(m: Mlp, x: [b D], step: int) -> [b C] = forward(m, x, Train(seed=step))

fn evaluate(m: Mlp, x: [b D]) -> [b C] = forward(m, x, Eval)

Mode is a parameter or return type, and a mode can be bound to a name, passed and returned like any value. A call that leaves out a Mode argument is E1201, and anything other than a mode where one is expected is E1202; a Python value fits, and is checked when the function runs (Runtime checks). A mode isn’t a tensor (E0412), and Mode in another type, a Mode parameter with a default, Train without (seed=…), and Eval(...) are E0303 or E0306. seed is an int, given by name (E0406; E0502 or E0805 for something other than an int). Mode, Train and Eval are built-in names, which a declaration or a value of the same name hides.

From Python, the modes are shapescript.Train(seed=…) and shapescript.Eval, both instances of shapescript.Mode. They are immutable.

dropout

dropout(x, p, mode) takes a tensor of any shape, a float p and a mode, and gives a float tensor of x’s shape and dtype whose domain is Real. With Eval it returns x itself. With Train(seed=…) each value is zeroed with probability p and the others are divided by 1 - p, so the mean is unchanged; gradients go through the values kept. p is at least 0 and less than 1, or dropout raises ValueError when it runs.

batch_norm

batch_norm(x, mode, over=C, weight=w, bias=b, running_mean=mu, running_var=var, eps=1e-5) is batch normalization, a stub in nn.ss: each index of the axis over, the channel, is normalized over every other axis of x, as torch’s BatchNorm1d and BatchNorm2d normalize over every axis but the channels. weight, bias, running_mean and running_var are [C] and may be left out. The mode decides which statistics it uses: with Train(seed=…), the batch’s mean and biased variance; with Eval, running_mean and running_var, which must be given then (ValueError otherwise). It draws nothing at random, and keeps no state: the running statistics are values like any other, such as a model’s buffers, and a training step updates them itself, with mean over the batch:

fn track(mut m: Net, x: [b D]):
    h = x @ m.w1
    mu = mean(h, over=b)
    m.var = 0.9 * m.var + 0.1 * mean(square(h - mu), over=b)
    m.mean = 0.9 * m.mean + 0.1 * mu

On several devices, its statistics with Train are the whole batch’s.

Random draws

The same seed gives the same random numbers, and two random operations never share a draw, even with the same mode:

• Paths. A training mode carries a path as well as its seed. Every call that passes a mode on, to a ShapeScript function, a prelude stub such as dropout, or a Python function, extends the path with the callee’s name and how many times this call of the enclosing function has already passed a mode to it: forward#0, then forward#1; dropout#0, then dropout#1. value_and_grad(f, m, …, mode) passes the mode on as the call f(m, …, mode) would, extended with f’s name, so a training step’s loss draws the masks that calling the loss function there would. Inside a loop, the count is per iteration, and the name says which iteration: forward@4#0 is the first call of forward in the iteration of for epoch in a..b: whose epoch is 4, and forward@4.17#0 the first in the 17th (from 0) iteration of a loop inside it over anything else, such as batches. So the two dropouts in one function, two calls of forward with the same mode, and each iteration of a loop that calls forward draw different masks, and a loop’s draws depend on its iterations, not on where it started: a loop over 3..5 draws what one over 0..5 draws in its last two iterations. A mode the checker sees only as a Python value isn’t extended, so give it its type, with a Mode parameter, before using it twice. A notebook cell counts its own calls outside any fn. Calling the same function from Python twice with the same mode gives the same masks both times.
• Generators. Each draw has a torch.Generator of its own on the declared device, seeded with the first 8 bytes of the SHA-256 of "{seed}/{path}", the path’s parts joined with /, as a model field’s seed is made (Compiled models). Nothing reads or changes PyTorch’s global random state. Masks are drawn as f32 whatever the input’s dtype, on the declared device, so a seed gives the same mask on the same kind of device; CPU, CUDA and MPS generators give different numbers.

Paths count calls, not source positions, so editing other functions, or adding a line before a call, doesn’t change its draws. Adding a call to the same function before it, in the same iteration, does.

Data and datasets

Data records

data Name: followed by an indented block of fields declares a data record: the fields of one example of a dataset.

## One labelled image.
data Digit:
    img:   [D]
    label: [] Index[C]

Each field is a name and a tensor type, and the record and its fields can have doc comments. A field’s shape is made only of numbers and dims, so every example has the same shapes and examples can be stacked into batches: a generic name is E0301, and ?k is E1302. A field can have any dtype and any domain. It can’t be a host scalar, a type without a shape, a model, a data or logs record, a dataset, a gradient, optimizer state, a Schedule, a Mode or a LogSink, and it can’t be named with a Python keyword or a name that starts with _ (E1302). A record’s fields have different names (E0302).

Values of different lengths, such as sentences, aren’t supported yet: every example of a record has the same shapes, and an example whose field has another shape is an error when it’s loaded. Pad them to a dim in the dataset and add a Mask field, as in tokens: [T] Index[V] and mask: [T] Mask. Padding varying axes for you is planned (FAST-3 in the plan).

Batches

[b] Digit is a batch of b examples, in which every field has the batch’s axes first: batch.img is [b D] and batch.label is [b] Index[C]. Digit alone is one example, whose fields have their declared types. A field the record doesn’t declare is E1301.

A batch type can be a parameter or return type, and in a parameter type, as with a tensor, a new name such as b is a generic axis that the arguments determine (Generic calls). Only a batch of the same record with the same number of batch axes fits; another record, a dataset or a tensor is E1303. A batch isn’t a tensor: arithmetic on one is E0412. A record’s name isn’t a value or a function (E0303), and a record’s name with a size, an over or a dtype, as in Digit[3], is E0303 too. as can’t give a value a batch type (E0303).

Datasets

Dataset[Digit] is a dataset of Digit examples. A dataset comes from Python: a parameter of that type takes any map-style dataset, an object with __len__ and __getitem__ such as a torch.utils.data.Dataset, whose items have the record’s fields, as a dict (other keys are ignored) or as a tuple or list in field order; for a record with one field, an item may be its value. In a notebook, x as Dataset[Digit] gives a Python value that type. Only a dataset of the same record fits a Dataset[Digit] parameter or return type (E1303). Dataset with anything but a record’s name in its brackets, or with a shape, is E0303.

Datasets are used with two pipe stages, which are built-in functions:

seed is given by name and n by position or as n= (E0406; E0502 or E0805 for something other than an int). Either stage on something that isn’t a dataset, shuffle after batches, batches twice, and a batch size known to be less than 1 are E1304. The order shuffle gives comes from a generator of its own, seeded from s, so it doesn’t depend on PyTorch’s global random state, and the same seed always gives the same order.

Looping over batches

for batch in train |> shuffle(seed=epoch) |> batches(64): runs its block once per batch. All batches but the last have n examples; the last has the examples that are left, so 100 examples in batches of 64 give a batch of 64 and one of 36. So the batch axis’s size isn’t known statically: the loop variable is [b] Digit, where b is an axis whose size is unknown, a new one for each loop, shown as ?b where it’s compared with other sizes. Inside the loop, b names that axis, as in batch.img as [b D], and is its size as a value, unless b already names something. Passing the batch to a function with a [b] Digit parameter makes the function’s b that size. The loader pads the last batch to the full size with masked-out examples, so that every batch has the same shape, when every operation across the batch has a masked version; b is then still the number of real examples where it’s used as a value (Fixed shapes and padding).

Loading

When a loop starts, the loader loads the examples in the loop’s order and checks each against the record as it’s loaded: the item must have every field, each field must be a tensor, a number, or anything torch.as_tensor takes, with the field’s shape exactly, and its values must be of the field’s kind. A float field takes any numbers, an integer field integers or booleans, and a bool field booleans; each is then converted to the field’s dtype. With value checks on (Runtime checks), each field is also checked against its domain. An item that doesn’t fit raises shapescript.ShapeError, or shapescript.DomainError for a domain, naming the .ss location of the field, the field and the item’s index in the dataset:

ShapeError: mnist.ss:8:5: field `img` of item 37 should be [D] = (784,), got shape (783,)

The examples of a batch are stacked field by field, and the batch is copied to the declared device. From Python, a batch is a shapescript.rt.Batch, whose fields are attributes. A function with a batch parameter also takes a dict with the record’s fields from Python, and checks each field as it checks a tensor argument.

Worker processes. Batches are built by worker processes, 4 unless the machine has fewer CPUs; shapescript.loader_workers(n) changes that, and 0 loads them in the calling process. The workers are started once per dataset parameter and kept for the loops over it. They are always started fresh, as with multiprocessing’s “spawn” start method, whatever the start method is, and never import the program’s main module: a worker unpickles the dataset, which imports the dataset’s own module, and nothing else. So the top level of a calling script runs once, and needs no if __name__ == "__main__":, even where Python’s default start method would import it again in each worker (“spawn” on macOS and Windows, and “forkserver” on Linux from Python 3.14). A dataset that refers to something defined in the main script, or that can’t be pickled, is loaded in the calling process instead, with a RuntimeWarning saying why. A module built with --standalone always loads in the calling process, since workers can’t import it.

Datasets kept on the device. A PyTorch torch.utils.data.TensorDataset with a tensor per field, in field order, whose item i is row i of each, isn’t loaded batch by batch. When a loop over it starts on the CPU, or on a GPU where it takes at most a quarter of the free memory, it’s checked once, as a whole, against the record, converted to the fields’ dtypes and copied to the device, and each batch is gathered there with an index operation: no loader workers, no copy from the CPU and no check per batch. It’s copied and checked again only if its tensors change (an in-place change to one counts). A subclass that changes __getitem__ or __len__, or adds __getitems__, as an augmentation does, is loaded as any dataset is; so is one with an item that doesn’t fit, which raises the error as the loader would. On several devices each worker keeps a copy of its own.

Logs and sinks

Logs declarations

logs Name: followed by an indented block of metrics declares a logs record, the numbers a program reports:

## How each value is combined when several are reported together, such as over an epoch.
logs StepLogs:
    loss:      [] mean
    grad_norm: [] max

Each metric is a name, a type, and how its values combine: mean (the arithmetic mean), sum, max, min or last (the latest value). A metric is a single number: its type is [], with any domain and dtype, as in [] NonNeg or [] i64, or a host scalar int, float or bool. Anything else, a missing or unknown combining word, and a name that is a Python keyword or starts with _, are E1402. A record’s metrics have different names (E0302). A logs record and each metric can have a doc comment.

Records

A logs record’s name is a type, as in -> StepLogs or r: StepLogs; another record or a tensor there is E1403. StepLogs(loss=l, grad_norm=n) makes a record. Every metric is given, by name: a metric the record doesn’t declare is E1401, and one given by position, twice or not at all is E0406. Each value must fit its metric’s type as an argument fits a parameter’s: a [] tensor or a number for a [] metric (E0403 for a tensor with axes, E0501 and E0502 for its domain and dtype), and a host scalar for a host one (E0805). r.loss reads a metric, with the metric’s type; a misspelled one is E1401.

A record holds its values as they are, so a tensor stays on the device until the record is written. It compiles to a subclass of shapescript.rt.Logs, and from Python its metrics are attributes.

Sinks

LogSink is the type of a sink, a Python object with a write(dict) method that records are written to. Sinks come from Python, as arguments, and ShapeScript code can only pass them on and call two methods on them:

• sink.write(record) copies the record’s values to the CPU, turns each into a Python number (float, int or bool, after the metric’s dtype), and calls the sink’s write with a shapescript.rt.LogRecord: a dict from each metric’s name to its number, in declaration order, whose name attribute is the record’s name and whose combine attribute maps each metric to how it combines. Writing anything but a logs record is E1403.
• sink.flush() calls the sink’s flush(), if it has one, and does nothing otherwise.

Any other method is E0301, and a wrong number of arguments is E0406. Something other than a sink where a LogSink is expected is E1404. A LogSink parameter checks at run time that its argument has a write method (Runtime checks).

Combining. The program writes a record whenever it likes, such as after every step, and a sink decides how often to report: it combines the records it has been given, metric by metric, by each metric’s rule. The sinks in shapescript.sinks take every=: with every=1, the default, each record is a row; with every=n, each row combines n records; and with every=None, a row is written only when the program calls sink.flush(), so flushing at the end of each epoch gives one row per epoch, with the epoch’s mean loss and largest gradient norm. A flush or close() writes a row for the records since the last row, if there are any. shapescript.sinks.combine(records) combines a list of records the same way, for sinks of one’s own.

• shapescript.sinks.csv(path, every=1) writes a CSV file: a header of the metrics’ names, then a row per combined group. The file is flushed after every row.
• shapescript.sinks.tensorboard(log_dir, every=1) writes each metric as a TensorBoard scalar tagged with its name, at the step of the last record in the row, counting records from 0. It needs the tensorboard package, and raises ImportError saying so without it.

Both are context managers, and close() flushes them.

Printing

print(a, b, ...) prints any values, separated by spaces, as Python’s print does: tensors on any device, host scalars, strings, models and logs records. It copies tensors to the CPU itself. A [] tensor is printed as a number, and a logs record as a dict of numbers. Its arguments are positional; a keyword is E0406. It has no result.

Checkpoints

save

save(path, values…) writes a checkpoint: the values, in order, to the file path, replacing it. The path is a string or a Python value, given first, by position (E1603); the values follow, by position (E0406 for a keyword or for no values). It has no result.

fn run(mut m: Mlp, mut opt: AdamW[Mlp], train: Dataset[Digit], first: int, epochs: int, mode: Mode):
    for epoch in first..epochs:
        for batch in train |> shuffle(seed=epoch) |> batches(64):
            l, g = value_and_grad(loss, m, batch, mode)
            update(mut m, mut opt, g)
        save("ckpt.pt", m, opt, epoch + 1)

fn resume(train: Dataset[Digit], epochs: int, mode: Mode) -> Mlp:
    m, opt, done = load("ckpt.pt", Mlp, AdamW[Mlp], int)
    run(mut m, mut opt, train, done, epochs, mode)
    return m

A checkpoint holds models (every field, params and buffers), optimizer state (everything state_dict() gives: the optimizer’s moments or momentum, its hyperparameters, its learning rate or schedule, and its update count) and host scalars, int, float and bool, such as an epoch or step count. Anything else is E1601. A tensor can’t be saved on its own: tensor state that training needs belongs in a model, as a buffer if it isn’t trained, so that its type, and the file’s, say what it is. A Python value is accepted and checked when save runs, which raises TypeError for anything a checkpoint can’t hold.

load

load(path, Types…) reads a checkpoint. After the path come the types of the values in the file, in order, written as in a signature: a model’s name (Mlp), optimizer state (AdamW[Mlp], SGD[Mlp]), or int, float or bool. Its result is a value of each type: one type gives one value, and more give results to name, as m, opt, done = load("ckpt.pt", Mlp, AdamW[Mlp], int); a different number of names is E0411. A value, a tensor type, or any other type, such as Grad[Mlp], a Schedule, a Mode or a logs record, is E1602, as is an expression that isn’t a type; a name that isn’t declared is E0301, and AdamW without a model, or of something other than a model, is E0303. The types are given by position (E0406), and at least one is (E0406).

What load makes is new: a model and optimizer state it loads belong to the function, which can pass them as mut (Changing models with mut).

Checkpoints at run time

save and load compile to rt.save(path, values…, where=…) and rt.load(path, types…, where=…), with where the call’s file:line:col; from Python they’re shapescript.save and shapescript.load, whose types are model classes, shapescript.AdamW[Mlp] and shapescript.SGD[Mlp], and int, float and bool.

The file is written with torch.save, to a file beside path that is then renamed to it, so a save that is interrupted leaves the last checkpoint whole. It’s a dict, {"format": "shapescript.checkpoint", "version": 1, "values": [...]}, with one entry per value:

• a model: {"kind": "model", "type": "Mlp", "fields": [...], "tensors": {"w1": tensor, ...}};
• optimizer state: {"kind": "optimizer", "type": "AdamW[Mlp]", "optimizer": "adamw", "model": "Mlp", "fields": [...], "of": 0, "state": {...}}, where state is the state’s state_dict(), fields are the params it was updating, and of is the position of the model saved with it whose params those were, or None;
• a host scalar: {"kind": "int", "value": 3}, and the same with "float" and "bool".

Each field is {"name": "w1", "kind": "param", "shape": [784, 256], "dtype": "float32"}.

Loading reads the file with torch.load(path, map_location=device, weights_only=True), so the file can’t run code, and every tensor goes straight onto the program’s declared device, never the device it was saved from: a checkpoint saved on a GPU loads on a CPU, and the other way round. Before it makes anything, load compares the file with the types, and raises shapescript.ShapeError (a ValueError) listing everything that differs: each field whose shape or dtype differs, with both, each field one has and the other doesn’t, and each value of another kind or type. Loading an Mlp saved with Hd = 128 as one with Hd = 256 gives:

ShapeError: train.ss:9:20: `ckpt.pt` doesn't match the types it's loaded as:
  value 1, `Mlp`: field `w1` is (784, 128) float32 in the file, but `Mlp` declares (784, 256) float32
  value 1, `Mlp`: field `w2` is (128, 10) float32 in the file, but `Mlp` declares (256, 10) float32
  value 2, `AdamW[Mlp]`: field `w1` is (784, 128) float32 in the file, but `Mlp` declares (784, 256) float32
  value 2, `AdamW[Mlp]`: field `w2` is (128, 10) float32 in the file, but `Mlp` declares (256, 10) float32

A file with another number of values is a ShapeError naming both lists (`ckpt.pt` holds 3 values (Mlp, AdamW[Mlp], int), but `load` was given 2 types (Mlp, AdamW[Mlp])), and so is a file save didn’t write. A model is compared by its name and fields, as a model parameter is, and an int in the file can be loaded as a float.

A model is made with the file’s tensors, without drawing new ones. Optimizer state is made with the file’s hyperparameters, learning rate and update count, so a schedule carries on where it was; saved with its model, it’s bound to the loaded model, and otherwise to stand-in params until update moves it to the model it’s given. Training that saves after n steps, loads, and carries on gives the same result as training that didn’t stop (tests/python/test_checkpoints.py checks it bit for bit, dropout and a schedule included). Random draws need the same keys in both runs for that, and they have them: a mode’s draws inside a loop are keyed by the iteration, so a resumed loop over done..epochs draws what the uninterrupted one draws in those epochs.

Modules and generic functions

Importing other files

Each .ss file is a module with its own namespace. import attention makes the declarations of attention.ss available as attention.name: its dims, domains, models, data and logs records, functions and extern fns, all with their full types. A qualified name can be used wherever the plain name could: attention.causal_attention(q, k, v), m: attention.Mlp, attention.Mlp.init(seed=0), [b attention.D], Dataset[attention.Digit], AdamW[attention.Mlp], attention.StepLogs(loss=l). A name the file doesn’t declare is E2304. Messages show another file’s types with its name, as attention.Mlp, and its shape aliases are written the same way, x: attention.Images.

Another file’s dims are named as this file names them wherever their axes go: the result of attention.f(x), or a field of an attention.Mlp, whose types say [b D] in attention.ss, is [b attention.D] here, so softmax(y, over=attention.D) finds the axis. Another file’s D is a different dim from this file’s D, even when both have the same size, just as two dims of one file are: [n C] + [n attention.C] is E0404 until an as renames one. If a file is imported under two names, its dims take the name of the import without as.

• import models.mlp imports models/mlp.ss and binds the last part, mlp, as import models.mlp as mlp would. Unlike import python, which binds the first part as Python does, an import always binds one name, so every qualified name has exactly one dot.
• import shapescript.std.fit as std binds std. The name an import binds is a module-scope name like any other: a fn fit and an import shapescript.std.fit in one file are E0302.
• There is no from attention import f: names of another file are always written with its name, so where each comes from is visible where it’s used.

The file is found next to the importing file, then in each folder of the Python path, as the import hook finds modules: sys.path when compiling from Python, $PYTHONPATH for shapescript check and build, and the workspace’s folders then $PYTHONPATH in the language server; then, from the command line and the language server, among installed packages. The standard library’s files, shapescript.std.*, are also built into the compiler, so they’re found without Python. No file is E2301. Files that import each other, directly or through others, are E2302.

An imported file is checked with the file that imports it, as one program; an error in it is reported once, as E2303 at the import that leads to it, with the errors as notes. Its warnings are left to checking that file itself.

Generated code. The generated module imports the other file through the import hook, _ss_mod_attention = _ss_rt.import_ss(__name__, __file__, "attention"), which looks next to the module’s own file, then on sys.path, and installs the hook if it isn’t; the name attention is bound to the module too, so Python sees main.attention. Each file’s Python is its own module, compiled once, so models and records of one file are the same classes everywhere. The hook’s cache entry for a file records the text of every file it imports, and which file each import found, so the importing files are compiled again when an imported file is edited, and when an import would now find another file: in a copy of the folder, with another sys.path, or where a new file hides the one it found. A --standalone build can’t import other files (E2305), and neither can a notebook cell (which can’t call generic functions either) or a program on several devices, whose workers run only the file’s own code, nor an exported function that calls another file’s (E2306). A batch loop doesn’t pad a batch it passes to another file’s function, which is compiled without knowing.

Overloaded functions

Several functions of a file may share a name when each takes a different model as its first parameter, so each model can have its own forward and loss. A call picks the one whose first parameter takes its first argument’s model, at compile time. It looks among the functions of that name the calling file sees and those of the file that declares the model, so forward(m, x) finds attention.forward for an attention.Mlp without the qualified name. Two with the same model are E2307, and so is sharing a name with a function whose first parameter isn’t a model, or with a generic function. Files are namespaces, though: when a call’s name means a generic function of the calling file, that’s the function, and the model’s file isn’t looked in. So the standard library’s fit calls its own run and train_step even if the model’s file has functions of those names, and finds only names it doesn’t declare, such as loss, next to the model. A library can add helpers without changing what its users’ code means; a call whose first argument is a model none of them takes, or not a model, is E2308.

Each one is its own Python function, _ss_forward__Mlp, named after its first parameter’s type as written, and the name forward is a function for calls from Python that calls the one for its first argument’s model class.

In the editor, typing m. after a value whose type is a model offers the model’s fields and every function whose first parameter takes that model: in the file, in the files it imports, and in the other .ss files of the workspace’s folders, even those it doesn’t import yet (but not those that import it, which would be a cycle), with generic functions whose first parameter is a type variable. A field completes as m.w1. A function rewrites m.forward as the call forward(m, …), with a tab stop for each parameter after the first that has no default, mut where the parameter is (train_step(mut m, mut opt, batch, mode)), qualified as attention.forward(m, …) when it’s declared in another file, and adds import attention after the file’s imports (or before its first declaration) when the file doesn’t import it yet. A generic function is marked when a function it calls on its first parameter, directly or through the file’s other generic functions, has no version for the model in the generic function’s file or the model’s, the functions a call would look in: calling it would be E2311.

Generic functions

A parameter type that is a single capital letter, optionally followed by digits, that nothing in scope declares (no dim, domain, model, record, logs record or built-in type) is a type variable: M in m: M, Grad[M], AdamW[M] and SGD[M] stands for a model, and E in [b] E, E and Dataset[E] for a data record. A function with a type variable is generic:

dim D = 8
dim C = 3

model A:
    param w: [D C] = normal(std=0.1)

data Ex:
    x: [D]
    y: [] Index[C]

fn loss(m: A, batch: [b] Ex) -> []:
    return cross_entropy(batch.x @ m.w, batch.y, over=C)

## One SGD step on any model that has a `loss`.
fn step(mut m: M, mut opt: SGD[M], batch: [b] E) -> []:
    l, g = value_and_grad(loss, m, batch)
    update(mut m, mut opt, g)
    return l

fn train(mut m: A, batch: [b] Ex) -> []:
    opt = sgd(m, lr=0.1)
    return step(mut m, mut opt, batch)

This is compile-time duck typing, as C++ templates are: there are no traits or interfaces to declare. A generic function isn’t checked on its own. Each call works out its type variables from the arguments (a parameter’s model, record, dataset or batch; two different types for one variable, or a variable no argument gives, is E2310) and checks an instance, a copy of the function with the types put in, so the function may do anything the types it’s called with allow. Its calls with a value of a type variable as the first argument, such as loss(m, …) and value_and_grad(loss, m, …), find functions as overloaded calls do, in the generic function’s file and the model’s, so a model brings its own forward and loss. A function it needs that the type doesn’t have, wherever the call is in its body, including in an if never taken, is E2311 at the call to the generic function, naming the function and where the generic function uses it; any other error in an instance is E2312 there, with the error as a note. A type variable that also names an axis of the signature, or one only the return type uses, is E2309. Instances that make other instances can go 16 levels deep (E2313).

The language server’s hover on a generic function lists its type variables and the functions it calls on them.

Generated code. Each instance is a Python function of its own in the module whose code calls it, as _ss_fit__Mlp__Digit, so mut, gradients, padding and placement are worked out for it as for any other function. The generic function’s name is a function for calls from Python that calls the instance for its first model argument’s class; one that no ShapeScript code calls has no instances, and calling it from Python raises TypeError saying so.

Packages

A ShapeScript package is an ordinary Python package whose folders have .ss files, and pip installs it like any other: pip install shapes_extra, from a wheel file, a local folder or a git+https:// URL. Its modules are imported with their dotted names, import shapes_extra.layers, which finds site-packages/shapes_extra/layers.ss and binds layers, with every declaration’s full types, as for a file of one’s own; an error in it is E2303 like any other imported file’s. docs/packages.md is the guide to making, building and publishing one; shapescript new NAME makes a package to start from.

Finding installed packages. An import looks next to the importing file, then on the search path, then in the standard library built into the compiler, then in the folders of a Python interpreter’s sys.path, where pip installs packages. From Python, the search path is already sys.path. The command line asks an interpreter: --python PYTHON, else the active virtual environment’s ($VIRTUAL_ENV), else $SHAPESCRIPT_PYTHON (which the shapescript command installed by pip sets to its own Python), else python3 (python on Windows). It asks only when an import isn’t found otherwise, and keeps the answer in the user’s cache folder (below) until a folder of the path changes, as installing a package with a .pth file does. The language server asks the interpreter it knows from its python option or shapescript/setPython (the VS Code extension sends the Python extension’s selected interpreter; see Python interop). An installed package’s files are read-only to the language server: opening one publishes no diagnostics, nothing is written beside it, and the VS Code extension opens .ss files under site-packages read-only (files.readonlyInclude).

Building. A package’s pyproject.toml names the build backend shapescript.build, which the shapescript package provides:

[build-system]
requires = ["shapescript>=0.4"]
build-backend = "shapescript.build"

It builds with hatchling (its [tool.hatch.build] options apply, and every file of the package’s folders goes in the wheel, .ss files included), checks every .ss file of the wheel with the others, failing the build with their errors, and adds to the wheel the Python module of each .ss file, a py.typed in each top-level package with .ss files, and a shapescript-build.toml in each folder with .ss files that names the compiler that built them. So Pylance, pyright and mypy know the installed package’s types, and from shapes_extra.layers import forward works without shapescript.install(). A hand-written .py file of a module’s name stays in the wheel instead, with W2501 as a warning. The sdist and editable installs (pip install -e .) are hatchling’s own: an editable install imports the .ss files from the source tree, beside which the editor writes the modules.

Running from site-packages. The import hook and the generated modules compile an installed .ss file like any other. Where __pycache__ beside it can’t be written, as in a read-only site-packages, the compiled Python goes to the user’s cache folder instead: $SHAPESCRIPT_CACHE_DIR, else $XDG_CACHE_HOME/shapescript, else ~/Library/Caches/shapescript on macOS, %LOCALAPPDATA%\shapescript\Cache on Windows and ~/.cache/shapescript elsewhere, under modules/, in a folder named by a hash of the file’s path. Generated modules of installed packages are never rewritten, since pip doesn’t keep the files’ times.

Versions. A package’s .ss files are written in the language of the compiler that built it, which shapescript-build.toml records (compiler = "0.4.1"). Patch releases don’t change the language, so a package needs at least the major and minor version that built it: an older compiler reports E2901 at the import instead of failing on what it doesn’t know, and so does the import hook when Python imports the module. A package’s own requirement on shapescript (dependencies = ["shapescript>=0.4"], as shapescript new writes) makes pip install a compiler new enough. A pre-release of the compiler writes itself instead, as in shapescript>=0.1.2b1, since pip takes pre-releases only for a requirement that names one.

Training loops

There are no hooks, callbacks or trainer objects: a training loop is a function the program writes, from value_and_grad, update, datasets, logs and checkpoints.

The README’s training example

examples/train_mlp.ss is the README’s training example, exactly (a test compares them): a two-layer classifier, train_step with clipping and logs, and fit(mut m, train, sink, epochs, mode), which makes AdamW state and runs epochs epochs, shuffling each with its number, writing a record per step, flushing the sink and saving m, the state and the number of epochs done after each. resume(train, sink, epochs, mode) loads that checkpoint and carries on up to epochs. Called from Python with mode=shapescript.Eval, so that dropout is off, its losses for 100 steps match a hand-written PyTorch loop’s (torch.optim.AdamW, clip_grad_norm_) to within 1e-5, on the CPU and on a GPU, and stopping after 5 epochs and resuming gives the 100 steps of a run that didn’t stop (tests/python/test_train_mlp.py). The mode is passed to every train_step, so with Train(seed=s) each step draws its own masks, and a resumed run draws new ones, as above.

The standard library’s fit

The standard library’s fit is a generic function, in shapescript/std/fit.ss in the Python package (and built into the compiler): fit(mut m: M, train: Dataset[E], sink, epochs, size, lr, mode) trains any model M on any data record E for which the program declares loss(m, batch, mode) -> [], next to the model; it makes AdamW state and runs run, which shuffles each epoch with its number, calls train_step (AdamW, the gradient’s norm clipped to 1, a StepLogs record of the loss and the gradient’s norm per step), flushes the sink and saves m, the state and the epochs done to ckpt.pt after each epoch. resume(like: M, train, sink, epochs, size, mode) carries on from that checkpoint; like only says which model to load. Import it and call it from ShapeScript:

import shapescript.std.fit as std

fn train(mut m: Mlp, data: Dataset[Digit], sink: LogSink, epochs: int, mode: Mode):
    std.fit(mut m, data, sink, epochs, 64, 3e-4, mode)

A model without the loss it needs is E2311 at the call. It has no hooks or callbacks, and it’s short enough to read on one screen; to change the loop, copy the file (shapescript.std.path("fit") gives its path). With size=64 and lr=3e-4 it gives the README’s fit’s losses on the README’s model, and the same as a fit written for a second model, with its own forward and loss (tests/python/test_generics.py).

Devices

Declaring devices

A program declares once where it runs. A .ss file declares it with a line at the top:

devices cuda[0]

The devices are cpu, mps (Apple GPUs) and cuda[...], with one or more different device numbers in the brackets, as in cuda[0] or cuda[0 1 2]. cpu[...], with two or more different numbers, is that many worker processes on the CPU, which stand in for GPUs: devices cpu[0 1 2] runs a program just as devices cuda[0 1 2] would, with three CPU processes in place of three GPUs, to test a program for several GPUs on a machine without them (Several devices). Anything else, including cpu[0], is E1001. Indented lines under the declaration are options, for running on several devices: shard params and split NAME, one per line (Options for several devices; anything else is E1002, and an option given twice E1903). With one device they do nothing, so a program keeps them when its devices line changes. A file has one devices line (E1003), before all of its declarations (E1004); comments and blank lines may come first.

From Python, shapescript.devices("cuda[0]") declares the devices, with the same text as the line. A notebook declares them with a devices line in a cell or from Python; a cell may repeat the devices an earlier cell declared, but not change them (E1003).

Devices at run time

There is no default device. A module with a devices line declares its devices when it is imported, and every ShapeScript function asks for the declared device when it is called, so calling one before anything has declared devices raises shapescript.DeviceError, which names both ways to declare them. Declaring the same devices again does nothing. Declaring different ones, whether from two modules or from a module and Python, raises DeviceError naming both declarations, and so does declaring a device the machine doesn’t have. With several devices, a module runs as Several devices says; a function compiled for one device, from a module without the same devices line, raises DeviceError when it’s called, rather than running on one device.

Every tensor ShapeScript code makes, such as causal[n m], a number that becomes a tensor, or a model’s field, is made on the declared device. A tensor argument must already be there: nothing is moved for you. With shape checks on (Runtime checks), each argument that is a tensor is checked, and one on another device raises DeviceError naming the argument and the device.

Building for a device

shapescript build --devices "cuda[0]" declares devices for a file without a devices line, as the line would; a file that has one must say the same (E1006). shapescript.compile(src, path, devices="cuda[0]") does the same from Python. A build with --standalone can’t ask the shapescript package which devices the program declared, so it needs a devices line or --devices (E1005), and it can’t be for several devices, whose runtime is the package’s (E1805). --devices "cuda[0 1 2]" compiles a file without a devices line for several devices, as the line would.

Several devices

With several devices, as devices cuda[0 1 2], the program runs once, in one process, and each tensor operation is split across the devices. Nothing in the program changes but the devices line: the same code runs on one device or several, and never asks which device it’s on. On N devices a program computes what it computes on one, up to the order in which floating-point numbers are added; where the compiler can’t keep that guarantee, the program doesn’t compile.

By default this is data parallelism (MGPU-1 and MGPU-2 in the plan): each device has a copy of every model and of the optimizer state, and a part of each batch. The options under devices change how the state is laid out (Options for several devices): shard params gives each device a piece of every param (MGPU-4), and split H splits every value with an axis of size H, such as attention heads, across the devices (MGPU-5).

The controller and the workers

The process that imports the module and runs the program is the controller. Each device has a worker process, which the controller starts when the program first needs the devices, and which runs only the tensor code the compiler made for it, never the program itself. Workers are started fresh, without the main module, as loader workers are (Loading), so the top level of the calling script runs once, and print, save, file writes and Python calls each happen once, in the controller.

The workers form a torch.distributed process group: NCCL with cuda[...], one GPU per worker, and Gloo with cpu[...]. That is the only difference between the two: cpu[...] is for testing, and for running a program meant for several GPUs on a machine without them; it isn’t faster than devices cpu, since the processes share one CPU.

Where functions run

The compiler works out where each fn runs, from what it does:

• In the controller, if it loops over a dataset, prints, saves or loads a checkpoint, writes to a sink, uses Python, takes a dataset or a sink, or calls a function that runs there. Such a function is ordinary Python in the controller; tensors it makes are on the controller’s CPU.
• On the workers, if it uses tensors, models, batches, gradients or optimizer state, and does none of those things. Every worker runs it, each on its part of the data. A call from the controller sends each worker its part of the arguments, and waits for them.
• Anywhere, if it does neither, as a function of ints does. It runs where it’s called.

Device work, value_and_grad, update and changing a model’s field, is on the workers, so it can’t be in a function that runs in the controller (E1803): it goes in a function of its own, as the README’s train_step is, which the training loop calls.

The batch axis

A batch is split across the workers along its batch axis, the first axis of a [b] Digit batch type, and so is every value the types give an axis of that size, such as batch.img, [b D], and x @ m.w1, [b Hd]. Each worker has some of the examples: a batch of n is split as evenly as possible, the first workers taking one more, so 64 over 3 workers is 22, 21 and 21, and 36 is 12 each. Nothing is padded or repeated. Values without the batch axis, such as a model’s params, are whole on every worker.

A function on the workers that is given a value with the batch axis has that parameter split too, so the split follows the data through calls: in the README’s example, forward’s x: [b D] is split because loss passes it batch.img. A function is compiled once, so it can’t take a split value in one call and a whole one in another (E1804). The batch axis stays an axis of its own: grouping it with another, as in (b n), making a value with two batch axes, a batch with a fixed size, or the batch’s size used as a number, which on each worker would be its part’s, is E1802.

Operations across the batch

An operation whose operand is split along the batch axis, but whose result isn’t, combines values across the workers. These have a distributed lowering, which gives what one device would compute, up to the order of the additions:

Each one is an all-reduce: every worker gets the whole result, so a loss or a logged value is the whole batch’s, and an uneven split counts each example once. Anything else that works across the batch axis is E1801: another reduction over it (max, min, std, var, logsumexp, norm, prod, argmax, any, all, reduce with anything but sum), an operation along it (softmax, cumsum, sort, topk, gather, scatter, an extern fn such as layer_norm given it as over=), indexing or slicing it, and causal with it. Operations along other axes, and everything per example, are unchanged.

Gradients and results

value_and_grad on the workers combines the gradients across them: each worker differentiates the value, which is the same on every worker, through its own part of the batch, and the gradients are added up, so they’re the gradient for the whole batch whatever the split. Then every worker applies the same update to its copy of the model, so the copies stay equal. A logs record holds whole values, so writing one to a sink in the controller writes the whole batch’s loss.

A call from the controller gives the first worker’s result, since every worker has the same; a result with the batch axis is every worker’s part, joined in order. Models and optimizer state stay on the workers.

Random draws differ. With Train(seed=…), dropout draws each worker’s mask for its part of the batch, so the masks aren’t one device’s, and results differ from one device’s beyond rounding; with Eval, they don’t. Tying draws to each example’s place in the dataset, so they match, is planned (Q-16).

Placement and communication on hover

With several devices, hovers in the editor show what the compiler planned (MGPU-6). The checker works out the table they show from the placement and the types (Checked::devices, a check::comms::DeviceTable): how each value is laid out, and every collective the compiled program runs, with its operation, its line, how many values it combines and of what dtype.

• A value says how it’s laid out: split across the 3 devices along b, each holding its share of the batch (22/21/21 for a batch of 64), with the example from the file’s batches(n); whole on every device; for a model, a copy on every device: 813,056 bytes (794 KiB) of params and buffers each, or with shard params sharded: each device holds a third of every param of Mlp, 271,024 bytes (265 KiB) of its params; with split H, split along H: each device holds 4 of H = 12; and in a function that runs in the controller, that it’s there. The batch axis’s own name says it’s split.
• A function says where it runs, and on the workers how many bytes it sends per call, its calls included.
• A line that sends anything, hovered anywhere, lists each collective on it, and each call of a function on the workers that sends something:

Sent between the 3 devices on this line, per call:

- all-reduce of the gradients of `Mlp`'s params `w1` and `w2`, which adds up every device's
  part of the batch: 203,264 × f32 = 813,056 bytes (794 KiB)
- inside `loss`: 16 bytes per call

In all: 813,072 bytes (794 KiB).

That is the README’s value_and_grad line on cuda[0 1 2]: the gradients are D·Hd + Hd·C values, and loss’s cross_entropy all-reduces a sum and a count, 8 bytes, forward and backward. The collectives are each lowering’s: the all-reduce of an operation across the batch (mean sends its result’s values and a count, cross_entropy 2 values, batch_norm 2·C + 1 in f64), value_and_grad‘s all-reduce of the whole params’ gradients, one per dtype; with shard params an all-gather of each param where it’s read, padded to a multiple of N, whose backward is a reduce-scatter, and norm and clip_norm‘s all-reduce of one number; with split H each gather (an all-gather) and each sum (an all-reduce) of the plan. A collective in a function that value_and_grad differentiates, of floats, counts twice, since the backward pass runs one as big. Sizes are numbers when the types give them and symbolic otherwise, with the dims’ values put in: 4·D·b·n = 3,072·b·n bytes. They’re the sizes of the values combined; an all-reduce on a ring sends about 2·(N−1)/N times that from each device.

Python calls on several devices

Python calls run once, in the controller, as they would on one device: a function that uses Python runs there (Where functions run). So a Python call can’t be given data that is split across the devices, a batch or a value with its batch axis, which would mean gathering the batch through the controller’s CPU at every step (E1901). That includes a batch in the controller’s own loop over batches(n), which goes to the devices in parts.

There are two ways round it. A function that prepares examples can run in the dataset, as it loads each example. Or it can be declared with extern fn: an extern fn works on each example by itself, along every axis it isn’t given as an Axis[...] argument, so it runs on the devices, each device’s worker calling it with its own examples:

extern fn random_crop(img: [b D]) -> [b D] = "augment.random_crop"

The workers import the module an extern fn names (here augment), from the controller’s sys.path, but never the script the controller runs, whose top level runs once: an extern fn naming a function in that script raises ImportError when the workers load the program. A random extern fn draws from each worker’s own random state, so, as with dropout, its draws differ from one device’s.

Options for several devices

Options under the devices line change how the program’s state is laid out over the devices. They change speed and memory use, never results, beyond the order of the additions:

devices cuda[0 1 2]
    shard params
    split H

shard params (MGPU-4) gives each device a piece of every param, about 1/N of it, and so of the optimizer state, as FSDP does: each param is flattened and cut into N equal pieces, the last padded. When the program reads a param, as in x @ m.w1, it’s gathered from the pieces (an all-gather), and its gradient, each device’s part of a sum, is added up and cut into the same pieces (a reduce-scatter), so each device updates only its piece. norm and clip_norm of a gradient add up the pieces’ squares across the devices. save collects the pieces into the file one device would write, and load cuts them again, so checkpoints move between one device, several, and several with shard params. Reading a field from Python, or a model’s state_dict(), gives it whole.

split H (MGPU-5) splits every value with an axis of size H, a dim, across the devices, each holding H / N of it: a model’s fields, the arguments from Python and every value in between. H must be a dim (E1904) that the number of devices divides (E1902), and one split is allowed (E1903). A value is split along its axis of size H, or along a group axis that starts with H, such as (H Dh), whose part on each device is a block of whole heads; a value with H in two places, or inside a group after another axis, is whole. Batches aren’t split with split H: every device has the whole batch.

The compiler carries the split through each operation:

A value that is computed whole but whose type has H is cut to each device’s heads, which needs no communication: in the README’s attention example, split(b n (3 H Dh) -> b H n Dh) cuts q, k and v to each device’s heads, the scores, masked_fill and softmax run on them, and the last matmul(w_o) multiplies each device’s block of (H Dh) by its rows of w_o and adds the parts up. Gradients follow the same plan backwards: a gathered value’s gradient is added up and cut, and a split param’s gradient is its device’s own. value_and_grad adds up the gradients of whole params, which every device computed from its own heads.

With shard params as well, the params with an axis of size H are split along it, and the others sharded.

Several devices from Python

From Python, a model on several devices is a handle, a shapescript.parallel.RemoteModel: its tensors stay on the workers, and reading a field, as model.w1, copies the first worker’s to the CPU. model.state_dict() copies them all. A handle can’t be changed from Python: a ShapeScript function that takes the model as mut changes it. Mlp.init(seed=…) makes the model on every worker, the same on each; adamw(m, …) and sgd(m, …) of a handle make optimizer state on every worker (RemoteOptimizer). save writes the first worker’s copies, in the same file format, and load makes the models and state on every worker; a checkpoint saved on several devices loads on one, and the other way round.

A function on the workers takes tensors and batches from Python as on one device, on any device: they’re copied to the workers, split as its parameters are. Errors raise in the controller as they would on one device, as a ShapeError with the .ss location, say, with the worker’s traceback as a note.

Speed on several devices

What a step costs on several devices, beyond the devices’ own work:

• Batches are loaded by the workers. A loop over a dataset’s batches in the controller, as the README’s run, doesn’t load them: the dataset is sent to every worker once, when the loop first needs the workers (again only if its length changes, so a dataset is taken not to change while it’s used), and each call sends a batch’s number. Each worker makes the loop’s order from its shuffles’ seeds, loads its own part, padded and split as the controller would split it (Padding on several devices), and checks its items as the loader does. It loads the next batches in threads of its own while the step runs, and on a GPU copies them there from pinned memory on a stream of its own. Before the step the workers agree that each has its part, so an item that doesn’t fit raises the loader’s ShapeError or DomainError in the controller and none of them starts the step; the models are as they were. A dataset defined in the main script can’t be sent, since workers never import it: the controller loads it, with a warning, and sends each worker its part, which costs a copy of every batch through the pipes. Reading a field of such a batch in the controller (it’s a shapescript.parallel.RemoteBatch) loads it there.
• A TensorDataset is kept on each GPU (Datasets kept on the device): each worker gathers its part of each batch from its own copy, the next batch while the controller makes the next call, and, since the whole dataset was checked when it was kept, doesn’t wait for the others to agree that each has its part. This is what DDP’s usual measurement does, with its data on the GPUs, and what makes a large batch as fast as DDP’s.
• A dataset can load a batch at once. If it has __getitems__(indices), PyTorch’s protocol for that, and returns a dict of stacked fields (each [n ...] for n indices), the batch is checked as a whole, which is much faster than item by item; a list of items is checked one by one. This holds on one device too.
• Tensors from Python are sent. A function on the workers called from Python with tensors gets them through the pipes, as their bytes: fine for small values, and costly for big batches (a 6144-example batch of 784 floats, 19 MB, takes about 20 ms a step). Loop over a dataset in ShapeScript instead.
• Steps are recorded on each worker (Recorded steps), with the gradients’ NCCL all-reduce in the graph, so a step is one launch per GPU.
• Each call is one small message per worker, the same for every worker when it can be, pickled once, and an answer from each; the first worker’s carries the step’s logs, which are read on the host every step when a sink writes them.

Measured on 3×A100-40GB with the README example (the plan’s MGPU-7 has the table), its data a TensorDataset: 1.37 ms a step at batch 64 and 1.68 ms at batch 6144, against 2.37 and 2.42 ms for DDP on the same machine, whose data is on the GPUs. With a dataset each worker loads from the CPU, 1.98 ms at 64 and 3.45 ms at 6144: copying 2048 examples per worker from the CPU every step is what costs. DDP’s step time differs between machines of the same kind (from 1.5 to 2.4 ms), more than ShapeScript’s.

When a worker fails

If a worker stops, as when it’s killed or runs out of memory, the controller raises shapescript.WorkerError straight away rather than waiting, naming the worker, how it stopped and what it was running. The other workers are stopped too, and what was on the devices is gone: new models, or a checkpoint, are needed to go on.

A worker waits for the others where they combine values. If they don’t all get there within the collective timeout, the program stops with a WorkerError that says where each worker was, by its .ss location:

WorkerError: a worker failed while running `train_step` (train.ss:32:4), and worker 1 didn't finish:
  worker 0: was waiting for the other workers at train.ss:30:12
  worker 1: didn't answer; it was idle
  worker 2: was waiting for the other workers at train.ss:30:12

The timeout is 600 seconds, or $SHAPESCRIPT_COLLECTIVE_TIMEOUT; shapescript.collective_timeout(seconds) changes it for workers started afterwards. Ctrl-C stops the workers along with the program.

Fixed shapes and padding

A training step whose shapes never change can be recorded once, as a CUDA graph, and replayed: one launch per step instead of one per operation (Recorded steps). Shapes change where a size varies: the last batch of an epoch is smaller, sequences have different lengths, and a boolean selection’s size depends on the data. ShapeScript fixes them by padding: the varying axis is padded up to a fixed size, the padded positions are masked out, and every operation across the axis is rewritten so that its result is exactly what it would be without padding, up to the order in which floating-point numbers are added. Where the checker can’t prove an operation exact, the axis isn’t padded and the program runs as it would otherwise. Padding never changes what a program computes.

Padding makes every step the same, which recording needs: with it, an epoch of the README’s training example replays one recorded graph.

What is padded

An axis is padded when its size is one symbol from one of three places:

• The last batch of a loop over batches (FAST-2). In for batch in train |> batches(64):, the loader pads a last batch of 36 to 64 with copies of its own examples, and every batch carries a mask of its real examples, all of them for a full batch. So every step of an epoch has the same shapes and does the same operations, the full batches included. This is on by default, wherever the checker finds it exact; where it doesn’t, the loop has hint H2001, which says why.
• The generic axes of a function that Python calls with varying sizes (FAST-3). shapescript.padded(f) wraps a ShapeScript function so that each call pads each generic axis it can, such as n in x: [n D], up to the next of a few fixed sizes: 1, 2, 3, 4, 5, 7, 9, 12, 15, 19, 24, 30, 38, 48, 60, 80, 104, 136, 176, 224, 280, 352, 440, 552, and so on, each about 25% above the last and a multiple of 8 above 64. Sequences of every length from 1 to 512 make 24 shapes. Its results are cut back to the real sizes, so the caller sees what f gives. shapescript.padded(f, "n") pads only n, and is a ValueError naming the operation if n can’t be padded.
• Boolean selections (FAST-4). x[mask] with a mask of one axis has an unknown size ?k. When every operation on the selection has a masked version, it’s kept at x’s size, with the mask saying which positions are selected, so its shape doesn’t depend on the data. A function that returns it, like the README’s nonempty_rows, gives it back kept so to a caller that can keep it too, and compacts it otherwise: Python always gets the compact value, as it would without padding.

Only these vary. A data field still has a fixed shape (E1302), so variable-length examples come into a program as a generic axis of a function called from Python, or are padded to a fixed maximum by the dataset, with a Mask field. Examples of different lengths in one batch (a mask per example) aren’t supported.

Operations across a padded axis

The checker follows a padded axis through the program by its size, as it follows the batch axis across several devices: every value whose type has an axis of that size is padded along it, including in the functions it’s passed to. Each operation across the axis is made exact:

Anything else across the axis has no masked version, and keeps the axis from being padded: sort, argsort, topk, argmax, argmin, gather, scatter and cat over it, cross_entropy over it, indexing or slicing it, a selection’s size used as a number (which would wait for the device), an as that gives it a fixed size, a pattern that groups it with another axis, print or a Python function given a padded value, an extern fn that isn’t elementwise, a parameter without a shape (unless the function only does elementwise maths), scaled_dot_product_attention’s keys, and, for a selection, whose positions aren’t in order, causal. A loop’s batch, or a function’s axis, is padded only if nothing it reaches is one of these; a selection that reaches one is compacted where it’s made.

Why padding is exact

The padded positions hold copies of real ones: the loader repeats the last batch’s own examples, shapescript.padded repeats a sequence’s positions, and a selection’s unselected positions are copies of its first selected one. So every value computed at a padded position is one a real position could have, as finite as a real one’s. Masking uses where, never a multiplication by the mask, so a padded value never reaches a result, and the gradient that reaches a padded position is exactly zero; since that position’s values are finite, the zero is never multiplied by an infinity, which would make a NaN. (Padding with zeros would break this: log of a zero is -inf, and its gradient times zero is NaN.) A masked reduction adds exact zeros and divides by the number of real positions, so the only difference from the unpadded program is the order of additions.

With Train(seed=…), dropout draws its masks for the padded shape, so the real examples’ masks may not be the ones an unpadded step would draw; with Eval there’s no difference.

When padding is on

Loops over batches and selections are padded wherever the checker proves it exact, and functions when wrapped with shapescript.padded. shapescript.fixed_shapes(False) turns all three off at run time, and shapescript.fixed_shapes(True) back on; it returns what it was. Off, the masked operations get no masks and are the plain ones, so the program runs exactly as it would have been compiled without padding. shapescript.shapes_seen() gives the distinct shapes each loop over batches and each padded function has seen, by the loop’s location or the function’s name, to check that an epoch has one shape. Notebook cells and exported programs aren’t padded.

The generated code shows what was rewritten: a loop over batches is _ss_rt.batch_loop(…, True, where), a function keeps the masks of its padded axes in _ss_pads, from its batches or from its caller (a keyword argument _ss_pad, which Python callers leave out), and a masked operation is a call such as _ss_rt.cross_entropy(z, y, 1, _ss_pads, ((4, 0),)).

Padding on several devices

A padded batch is split across the devices like any other, and its mask with it: 64 padded examples are split 22/21/21, and a device may have only padding. The distributed lowerings of a padded batch add up each device’s masked sum and its number of real examples, so a mean, or the batch’s cross_entropy, is over the real examples of the whole batch.

Recorded steps

A step is a function that changes a model or optimizer state, a mut parameter of either, and does device work itself: value_and_grad, update or changing a field. The README’s train_step is one; run, which calls it in a loop, isn’t, since it does no device work itself. The checker records whether each step can be recorded: it can unless it, or a function it calls, loops over a dataset, prints, saves, loads, writes to a log sink, uses Python or calls a Python extern fn (those for torch.… and the prelude’s are device work), waits for the device with if on a tensor, or makes a size that depends on the data that isn’t kept padded. A step that can’t be recorded has hint H3101 naming what prevents it, and runs operation by operation. Hovering a step’s name says whether it’s recorded, and hovering what keeps it from being recorded says so.

A recordable step is compiled as train_step = _ss_rt.recorded(train_step, "train_step", where). On CUDA, when it’s called:

• The first call runs as it is, on a side stream: the warm-up CUDA graphs need, in which the optimizer makes its state and cuBLAS its workspace. Then the step is recorded for that call’s signature.
• A call with a new signature records the step as a CUDA graph: forward, backward, with several devices the gradients’ all-reduce (NCCL collectives can be recorded), and the optimizer’s update. The signature is the shapes, dtypes and devices of the tensors it’s given (batch fields and their pads included), which model and optimizer state (and their tensors), its host values (int, float, bool, strings), and whether each mode is Train or Eval. So with padding an epoch is one graph; without it, the last batch’s shape is a second.
• Later calls copy their tensors into the graph’s own inputs, replay it, and return copies of its results.

What a step does on the host happens around the replay: the learning rate its optimizer’s schedule gives for its update count is written into the device tensor the update reads (AdamW is made with capturable=True on CUDA), and the count goes up after. Each dropout draw of a recorded step has a generator of its own, which the graph reads its seed and offset from on each replay; before each replay it’s seeded with the seed the draw would have had without recording, from the mode’s seed and path (Random draws). So a replayed step draws different masks at every step, and the very masks it would have drawn without recording: results with and without recording are the same, bit for bit on the RTX 5090 it was measured on. A step that reads a padded batch’s real size as a number is recorded for each real size.

If recording fails at run time, as when a step calls Python through a value the checker can’t see into, the step runs operation by operation from then on, with a RuntimeWarning. A step records at most 64 signatures; further ones run as they are. Recording only happens on CUDA: on the CPU, MPS and CPU workers, steps run as they are.

shapescript.record_steps(False) turns recording off, and shapescript.record_steps(True) back on; it returns what it was. It’s on by default, unless $SHAPESCRIPT_RECORD_STEPS is 0, false, off or no. shapescript.recorded_steps() says, for each step by name, how many graphs it recorded and for which shapes, how many calls replayed, ran as they are, or waited for a recording, and how long recording took; with several devices, the first worker’s.

Recording ahead

A loop over batches(n) knows its epoch’s batch shapes before it loads them, from the dataset’s length, the batch size and padding, and announces them when it starts. A step that has run once is recorded for each announced shape it hasn’t seen, in a thread of its own on a stream of its own, while the steps run, so the first step with that shape replays a graph that’s ready instead of recording one. A step announced before its first call is recorded for those shapes after its warm-up. While a recording is under way in the background, replays go on; a step that would run operation by operation waits for it, since a step run as it is beside a recording can spoil it, and recorded_steps() counts such waits.

With several devices each worker is told its part of each announced shape with the next call, and records it after answering, until the controller’s next message comes: NCCL takes collectives from one thread at a time, so a worker doesn’t record in the background.

In the editor

Hovering a loop’s batches says whether the last batch is padded, and if not, which operation prevents it and where. Hovering that operation says what it keeps from being padded. Hovering a function’s name says which of its generic axes shapescript.padded can pad, and hovering a selection whether it’s kept padded.

Python interop

import python torch binds the Python module torch. Members of a Python module (torch.load), calls to them, and anything computed from a Python value are Python values. A Python value has no type: the checker knows nothing about it, not even whether it is a tensor, and arithmetic or a comparison with one gives another Python value. A Python value can be:

• passed back to Python, or combined with other values;
• passed for a shapeless parameter or a host scalar parameter, which takes it unchecked, or for a model, Mode, optimizer state or Schedule parameter, which checks it at run time; it can be a learning rate or a hyperparameter of an optimizer, the model or state of update, the gradient of clip_norm, and the path or a value of save or the path of load, which save checks when it runs, but not the model of sgd or adamw, whose state’s type comes from the model’s (E1501);
• an if condition, an end of a range, or what a for loops over, whose items are Python values too (Control flow);
• given a type with as: a tensor type such as as [D 3*D], or as int, as float or as bool (Host scalars).

import python torch

fn load_qkv() -> [D 3*D]:
    return torch.load("w_qkv.pt") as [D 3*D]

Using a Python value where a shape is needed, such as returning it from a function with a tensor return type, passing it as a tensor argument, giving it to a pattern operation or to an operation along a named axis, or converting it with to and a domain, is E0413.

In the other direction, every ShapeScript function checks its arguments against its signature when it is called, whether from Python or from ShapeScript (Runtime checks).

In the editor, go to definition on the module of an import python (on any part of torch.nn, up to that part), on a member of one (torch.load, augment.random_crop), or on an extern fn’s path string opens the Python source. The language server asks a Python interpreter: the python initialization option or the last shapescript/setPython notification’s {python} (the VS Code extension sends the Python extension’s selected interpreter, else the shapescript.debug.python setting), else python3 (python on Windows). It runs a small script that finds the module with importlib.util.find_spec, looking in the file’s folder and the workspace’s folders first, then finds the name in the module’s source with Python’s ast, without running the module; a name the module only imports, as torch does load, is found by importing the module and asking inspect, and one defined in C by the package’s .pyi stubs. Answers are kept for each interpreter and name until the server restarts.

The other way, the language server answers the request shapescript/declarationOf with {path, name}, the location of the top-level declaration name in the .ss file path (the first part of a dotted name such as Mlp.init; an empty name is the file itself), which the VS Code extension uses for go to definition in Python files (Python modules for .ss files).

Python modules for .ss files

Python can import a .ss file in two ways: with the import hook, after shapescript.install(), or through the Python module generated beside it, model.py for model.ss. The module is what Python tools read: Pylance, pyright and mypy know nothing of .ss files, so without it from model import f is an unresolved import to them, even when it runs. It has:

• a first line that says it’s generated: # Generated by ShapeScript from model.ss. Edit model.ss; this file is rewritten when it changes.;
• a typed declaration of everything the compiled file defines, each with a docstring holding its ShapeScript declaration and its ## comment, which editors show on hover;
• a last line, __import__("shapescript.pymodules").pymodules.load(globals(), __file__), which compiles model.ss, with the import hook’s cache in __pycache__ and the check modes given to install() if it was called, and puts the compiled dims, models and functions in place of the declarations. So import model works without shapescript.install(). As with the hook, the code has the .ss file’s lines, so tracebacks and debuggers show ShapeScript, and a file with errors raises ShapeScriptError from its first error’s line (see Debugging).

The module imports torch and typing as _ss_torch and _ss_typing, so the table’s torch.Tensor is written _ss_torch.Tensor: that way the only names the module offers, as after from model import in an editor, are the file’s own (editors show the types themselves, Tensor). Tensor types become torch.Tensor; host scalars int, float or bool; a model, its class; Mode, Grad[M], optimizer state, Schedule, Dataset[R] and batches, the classes of shapescript.rt; and a LogSink, typing.Any. Parameters keep their order and defaults, as in the compiled function. Shape aliases are only for types, so they aren’t declared.

A file that imports others is checked with them, found next to it and then on sys.path (from Python) or $PYTHONPATH (the command line). Its module doesn’t import theirs, which may not have been generated: another file’s model is torch.nn.Module, its data record’s batches "_ss_rt.Batch" and its logs record "_ss_rt.Logs". When Python imports an imported file through its own generated module first, the importing file uses that module rather than compiling the file again, so models and records are the same classes. Overloaded functions are one declaration of the name, forward: _ss_overloads_forward, typed by a typing.Protocol whose __call__ has each as a typing.overload with its ShapeScript signature, so that editors find one place for the name (the variable’s docstring lists every signature); overloads that Python’s types can’t tell apart, such as two that take models of other files, both torch.nn.Module, are declared once, with every ShapeScript signature in the docstring, since type checkers reject an overload that can never be picked. A generic function is declared once under its name, which is the function that calls the instance for its argument’s class: a type variable’s model is torch.nn.Module, and its record’s batches "_ss_rt.Batch".

With the hook installed, the hook imports model.ss itself (it comes before Python’s own finder), so the module’s declarations are only read by tools, and the file is compiled once either way.

The language server writes the module when a .ss file is opened or saved and has no errors; the VS Code setting shapescript.pythonModules turns this off. shapescript pymodules FILE... writes them from the command line, and shapescript pymodules --check FILE... only fails if one is missing or out of date, for CI. The import hook rewrites a generated module that is older than its file. None of them ever replaces a .py file that doesn’t start with the generated line: that is someone’s own module, and the language server and pymodules report W2501 instead.

Going between Python and ShapeScript. In VS Code, Ctrl-click (go to definition) on a name a Python file imports from a .ss module goes to its declaration in the .ss file, never to the generated module, and so does Ctrl-click on a declaration in the generated module itself. Pylance, as any Python language server, answers with the generated model.py, and VS Code shows every extension’s answers, with no way for one to drop another’s; so the extension adds no answer of its own when another one answers, and when an editor then lands where that answer said, it shows the .ss file at the same declaration in its place, in the same tab. Going back from there skips the generated module. Without a Python language server, the extension’s own answer is the .ss file. A generated module opened on purpose, from the explorer or its tab, stays open: only an editor that lands where a go to definition just said is replaced. The module is looked for as Python would find it, next to the Python file, then in the workspace’s folders; a name of a module that is only a .ss file (which the import hook imports) goes to it too. The other way, the references of a .ss file’s top-level declaration (Find All References, and what Ctrl-click on the declaration shows) include the Python code that uses it: the uses Pylance finds of the generated module’s declaration, leaving out the generated modules, or, without a Python language server, the uses in the workspace’s Python files that import the module, read from their imports (from model import f as g, import model as m) and text, not their scopes. Peek Definition still shows Pylance’s answer, the generated module, as a peek doesn’t open an editor.

The modules are derived files. A project can ignore them (*.py beside .ss files in .gitignore), and they’re written again whenever the editor opens a file; but a type check in CI then needs shapescript pymodules first. Or a project can commit them, so that import model type-checks for anyone, with shapescript pymodules --check in CI to keep them up to date. Either way works; ignoring them keeps diffs to the .ss files.

Runtime checks

shapescript build compiles a file to a Python module. Every function checks its arguments when it is called, from Python or from ShapeScript: each parameter with a shape gets a rank check and a check of its sizes, and the function’s generic sizes are read from the arguments’ shapes (the body may use them, as in causal[n m] or sqrt(Dh)). One read from an axis that isn’t just it, as h is read from (h Dh), 16h or h+3, is checked first: some whole h of at least 0 must give the axis’s size, as for a call the checker sees, and the error says which sizes fit, as in whose axis 3 is a multiple of Dh = 16, got shape (2, 5, 72): 72 isn't divisible by 16. Every as with a shape compiles to a shape check, and so does as int, as float or as bool on a Python value. A parameter whose type is a model checks that its argument is a model of that name whose fields have the declared shapes and dtypes, and one whose type is Grad[M] that its argument is a gradient of M whose fields have the params’ shapes and dtypes. A parameter whose type is a batch of a data record checks each field as a parameter of the field’s batched type would be checked, and a dataset parameter wraps its argument so each item is checked as it’s loaded (Loading). A parameter whose type is optimizer state, such as AdamW[Mlp], checks that its argument is state of that optimizer for a model of that name with the same params, and a Schedule parameter that its argument is a schedule or a number (TypeError if it isn’t). A Mode parameter checks that its argument is a mode, raising TypeError if it isn’t. A parameter whose type is a logs record checks that its argument is a record of that name, and a LogSink parameter that its argument has a write method. Each argument that is a tensor is also checked to be on the declared device. Parameters without a shape, and host scalar parameters, aren’t checked.

Each as with a domain, and each parameter with a shape and a domain that has a test, also gets a value check. A value check tests what the domain promises: that values are within a range’s finite ends; that Prob over C sums to 1 over C, as does the exp of LogProb over C (to within the square root of the dtype’s epsilon); that Binary holds only 0 and 1, Additive only 0 and $-\infty$, and Index[V] integers in $[0, V)$; and that a Mask is stored as bool. Real and Logit have no check, and other dtypes aren’t checked.

debug skips the checks when Python runs with -O. A failing shape check, host check or model check raises shapescript.ShapeError, a failing device check shapescript.DeviceError, and a failing value check shapescript.DomainError. They name the .ss location; ShapeError and DomainError are ValueErrors, and DeviceError is a RuntimeError. shapescript run (python -m shapescript) shows these errors, and a .ss file’s compile errors, as shapescript check shows its diagnostics, without a traceback: the line of the .ss file and the call that ran it, and for an argument, the argument in the call with the parameter it was given for. Every other exception keeps its traceback, and so do these with SHAPESCRIPT_TRACEBACK=1 and under a debugger, which stops where the error is raised. On a CUDA tensor, value checks use torch._assert_async, so a failure shows up later as a device-side assert.

Export

shapescript export turns a function into a program for torch.export: a .pt2 file that torch.export.load reads and runs with torch alone, with no ShapeScript and none of the program’s Python, or an AOTInductor package compiled ahead of time, which C++ can load too. torch.export traces a function with example inputs and records the torch operations it runs; its hard part is saying which input sizes may change and how they relate. ShapeScript’s types already say that, so the example inputs and the dynamic shapes come from the function’s signature, and nothing is written by hand.

Export functions

export fn marks a function that is meant to be exported. The checker then checks it, and every function it calls, at any depth, for what can’t be traced into torch operations; shapescript check and the editor report these errors as for any other. An unmarked function can be exported too, with shapescript export --fn name, which checks it the same way first. A function that something exported calls isn’t marked itself: the errors in it name the exported function that reaches it.

In an exported function, and everything it calls:

• Python is an error. A name from import python, or anything reached through one, such as torch.load(...), torch.pi or x |> torch.tanh, is E1701: torch.export can only record torch operations, and a Python call could do anything. In a notebook, so is a Python variable. Give a torch function a type with extern fn instead.
• An extern fn must name a function in torch. Its path must be in the module torch or one of its submodules, or in math, which works on numbers and sizes; the prelude’s stubs into shapescript.rt, such as layer_norm, are torch code too. Any other module is E1702.
• as on a Python value is an error (E1703), since such a value already is one.
• Host-only values and operations are errors (E1704): models, modes, logs records, sinks, gradients, optimizer state, schedules, datasets and batches, the built-ins that make or use them (print, save, load, value_and_grad, update, sgd, adamw, clip_norm, cosine, warmup, shuffle, batches, Train), changing a model’s field, and a for loop over anything but a range. An exported program is for inference: training happens before it.
• return inside an if on a tensor is an error (E1706), as Control flow explains.

The parameters of the function being exported are the program’s inputs. Each must have a tensor type with a shape; an int, float or bool parameter with a default is fixed at the default, and anything else is E1705.

Dynamic shapes

Each input’s example tensor and the dynamic_shapes argument of torch.export.export come from the parameters’ types:

• A literal size, or one made of dims, such as 3*D, is static.
• Each generic axis, such as b in x: [b n D], is a torch.export.Dim, the same one wherever the axis appears, so [n D] and [n n] Mask share n.
• A size that is linear in one generic axis is a derived dim: 16gh and (gh 16) are 16 * gh, and n+1 is n + 1. Since every axis has at least one element, each Dim’s least value makes every size derived from it at least 1.
• Any other size, such as (h w), is Dim.AUTO, and torch works out what it can.

Each generic axis gets an example size different from every other size in the signature, so that torch can’t mistake one axis for another. Examples have the parameters’ dtypes and values in their domains, and are made on the program’s device: the file’s devices, or --devices, or the CPU. The .pt2 file doesn’t keep them.

Torch assumes that a dynamic size isn’t 1 wherever an operation would treat 1 specially, as reshape does, and the program would then refuse inputs where it is. shapescript export traces with torch’s backed_size_oblivious setting first, which reasons about 1 as about any other size, and without it if that fails, so generic axes can be 1 where torch allows it.

Assertions in exported programs

While torch.export traces, dynamic sizes are symbols, not numbers, and no values are known, so checks can’t run in Python. Instead:

• Each shape check of an as is a torch._check per axis. torch.export takes it as a fact, so rows as [196 D] makes the rest of the function see 196 rows, and the program checks it when it runs, raising RuntimeError if it fails. --checks decides which are kept, as for build.
• Each unknown size (?k) is a data-dependent size in the program, with what the checker knows of it, that it’s at least 0, passed to torch._check.
• Each value check, of an as or an argument, is an aten._assert_async in the program, which raises RuntimeError with the check’s message when it fails. --value-checks decides which are kept, and debug keeps them unless the Python that exports runs with -O.
• Argument checks become facts too; the program also checks its inputs’ shapes against the dynamic shapes itself.

So nonempty_rows(tokens[0]) as [196 D] exports, and the program raises an error when an image has an empty patch.

Control flow in exported programs

An if in an exported function becomes a function for each branch, of the values from before the if that the branches may assign, giving those back with the values both branches bind. If the condition is a bool known while the function is traced, such as one made of dims, Python runs one branch and the program has no trace of the other. Otherwise, for a [] Mask or a bool that depends on a dynamic size, the if is a torch.cond, and both branches are in the program. A branch of a torch.cond gives values rather than returning from the function, so return inside an if on a tensor is E1706; a return inside an if on a bool makes it a Python if, which needs the bool to be known while tracing.

A for i in a..b becomes a function of i and the values from before the loop that it may assign, giving those back. If both ends are numbers while tracing, as with 0..3 or 0..D, the loop is unrolled. Otherwise, as with 0..n for a generic n, it’s a torch.while_loop. A loop with a return inside is a Python loop, which needs its ends known while tracing.

In a torch.cond or a torch.while_loop, an int, float or bool that the block may assign, and a loop’s i, are 0-d tensors, since a program can’t keep a number that changes; a counter in a loop over 0..n works.

shapescript export

shapescript export model.ss --fn causal_attention -o attn.pt2
shapescript export model.ss --fn causal_attention -o attn.pt2 --aoti
shapescript export model.ss --fn forward --hf model_hf [--vllm]
shapescript export model.ss --fn forward --gguf model.gguf [--arch llama]

The last two export a model rather than a function, for Hugging Face and for serving engines. Both take --weights FILE, --seed N, --dtype float32|bfloat16|float16, --tokenizer DIR and --max-len N, and --hf takes --no-verify.

--fn names the function, and defaults to the file’s only export fn. -o defaults to the input with a .pt2 extension. --aoti writes an AOTInductor package instead of the torch.export program, with torch._inductor.aoti_compile_and_package, which needs a C++ compiler; torch._inductor.aoti_load_package loads it. --checks, --value-checks and --devices are as for build.

Tracing runs in Python, which needs torch and einops but not ShapeScript: the compiler generates a module that holds the functions, the runtime helpers they use and the export step, and runs it with --python, or $SHAPESCRIPT_PYTHON, or python3 (python on Windows). The shapescript command that a pip install provides sets SHAPESCRIPT_PYTHON to its own Python. --emit-python FILE writes the module instead; python FILE -o attn.pt2 [--aoti] runs it. A function that doesn’t exist, or a missing --fn when the file doesn’t have exactly one export fn, is an error of the command, with exit code 2; an error from torch while exporting has exit code 1.

The program’s inputs are the tensor parameters, in order. It runs with torch alone:

import torch
attention = torch.export.load("attn.pt2").module()
y = attention(x, w_qkv, w_o)       # any b and n

Exporting from Python

shapescript.export(path, fn=None, output=None, *, aoti=False, checks="always", value_checks="debug", devices=None) does the same in the calling Python and returns the torch.export.ExportedProgram. With output, it writes the .pt2 file, or the AOTInductor package with aoti=True. A file with errors raises shapescript.ShapeScriptError, and a missing function ValueError.

Hugging Face models

shapescript export asr.ss --fn transcribe --hf asr_hf --weights ckpt.pt

shapescript export --hf DIR turns a model into a Hugging Face folder, which transformers loads with trust_remote_code=True and which can be pushed to the Hub. Any model can be, not only a language model: transformers runs any PyTorch module. The function --fn (by default forward) must take the model first (E3306), and is checked as an exported function is, except that the model may be passed around and its fields read; its other parameters are tensors with shapes, or have defaults (E1705). The folder’s ShapeScriptModel holds the model’s fields as its parameters, under their own names, and its forward is the function, compiled: it takes the function’s other parameters, by position or by name, and gives what the function gives.

from transformers import AutoModel
asr = AutoModel.from_pretrained("asr_hf", trust_remote_code=True)
logits = asr(audio)                # transcribe(m, audio), for any b and t

The folder holds config.json, model.safetensors and modeling_shapescript.py, which needs torch, einops and transformers, not ShapeScript. After writing it, the export loads it with transformers, if that’s installed, and checks that it gives what the function gives, on example arguments made from its types, as for torch.export. The options are those of the next section.

If the function is a language model’s that serving engines can run, the folder is one that vLLM serves, as the next section says. If it isn’t, the export says why in a note, and config.json says so too (shapescript.vllm is false, and shapescript.unservable gives the reason); --vllm makes the serving check’s errors errors.

Serving engines

shapescript export --hf turns a language model into a Hugging Face folder that vLLM serves, and --gguf into a GGUF file for llama.cpp. A serving engine doesn’t run a model the way training does: it packs the tokens of many requests along one sequence axis, one request after another, keeps the keys and values of past tokens in a cache of its own, and computes logits only for the tokens it samples. So it runs three parts of the model itself, the token embedding, attention and the output layer, and the rest of the model as written. Usually each model is written again by hand for each engine. Here the rest of the model is the source, compiled, and the checker proves that the engine may run it that way.

The forward function, --fn, by default forward, must have this form:

fn forward(m: Llama, ids: [b n] Index[V], pos: [b n] Index[P]) -> [b n V] Logit over V:
    x = embedding(ids, m.embed)
    for l in 0..L:
        x = block(m, x, pos, l)
    return x |> rms_norm(over=D, weight=m.norm, eps=1e-5) |> linear(m.output)

These are the rules. With --hf, a function that breaks one is exported for transformers only, with a note; with --vllm or --gguf, it’s the error each rule names.

• Its parameters are the model, the token ids, [b n] Index[V] with V a dim, and optionally their positions, [b n] of an integer dtype. Other parameters need defaults. Anything else is E3301.
• It ends with the output layer: one return of linear(h, m.W) or h |> linear(m.W), with m.W a param of shape [V D] and no bias, which the engine runs (E3302). m.W may be the embedding table, which ties the two.
• The token ids are used once, by embedding(ids, m.E) with m.E a param of shape [V D], which the engine runs; the embedding table and the output layer are read nowhere else (E3303).
• Everything but attention treats each position on its own. Along n, only scaled_dot_product_attention with attn_mask=causal[n n] written in the call may work across positions; the engine puts its own attention in its place, which keeps the requests apart and reads its cache. A sum, a softmax or a contraction across n mixes requests, even where a mask would make padding exact, so attention written by hand, attention with another mask, and positions computed from n are E3304. The check follows n into every function the forward function passes it to, as padding does. Positions come in as a parameter.
• Attention takes [b H n Dh]: queries, keys and values of one shape, with H and Dh dims, the same for every attention, and at least one attention (E3305). Keys and values with fewer heads are repeated to H before attention, as repeat(b Hkv n Dh -> b (Hkv R) n Dh) does.

Everything else that an exported function can’t do, a serving export can’t either, except that the model may be passed around and its fields read.

The engine runs the attentions in the order the forward function runs them, each with its own part of the cache, so the forward function must run the same attentions on every call; for l in 0..L over a dim does. The checks of devices, dtypes, models and argument values are left out, since engines give the model its inputs on their own device and often in a lower precision; shape checks stay.

Hugging Face folders and vLLM

shapescript export llama.ss --hf llama_hf --weights ckpt.pt
vllm serve llama_hf --trust-remote-code --model-impl transformers --skip-tokenizer-init

The folder holds config.json, generation_config.json, model.safetensors and modeling_shapescript.py: the compiled forward function, the runtime helpers it uses, and the classes ShapeScriptConfig, ShapeScriptModel (the model up to the output layer) and ShapeScriptForCausalLM, which transformers loads with trust_remote_code=True. It needs torch, einops and transformers, not ShapeScript. The embedding table is model.embed_tokens.weight, the output layer lm_head.weight unless it’s tied, and every other field model.<field>, with its shape as declared, stacked layers included. num_hidden_layers is the number of attentions the forward function runs, counted by running it once.

Each attention calls the attention function the model is set to use, from transformers’ table: by default eager, which is the source’s attention, so transformers gives the source’s logits; vLLM’s Transformers backend sets its own, which keeps each layer’s keys and values in its paged cache, and runs everything else through torch.compile and CUDA graphs. vLLM’s own embedding and output layer replace embed_tokens and lm_head. transformers keeps no cache of past keys and values for the model (use_cache is false), so generate there computes each step from the start; vLLM keeps its own. Keys and values with fewer heads than queries are given to the engine repeated, so its cache holds H heads.

An engine reads each attention’s scale once, when it builds the model: ShapeScriptAttention.scaling is the scale the source gives, recorded in config.json as attention_scales when the source gives one, or torch’s default, $1/\sqrt{D_h}$; a scale that differs on a later call is an error. In a model served in bfloat16 or float16, what the source computes in f32 on purpose stays f32, as rotary angles from pos to f32 usually are, and attention and the output layer are given the model’s dtype, as the engine runs them.

After writing the folder, the export loads it with transformers, if that’s installed, and checks that it gives the logits the source computes, to 1e-4 of their size in float32 (5e-2 in the other dtypes); --no-verify (verify=False) skips this.

Without --weights, the weights are new ones made from --seed (by default 0), as M.init(seed=…) makes them; --weights takes a checkpoint that save wrote, whose first model of the forward function’s type is used. --dtype is float32 (the default), bfloat16 or float16. --tokenizer DIR copies a Hugging Face tokenizer into the folder; without one, the engine takes and gives token ids, as --skip-tokenizer-init tells vLLM. --max-len sets max_position_embeddings, by default P of the positions’ Index[P], or 4096. Tensor and pipeline parallelism aren’t supported yet.

GGUF files and llama.cpp

shapescript export llama.ss --gguf llama.gguf --weights ckpt.pt
llama-server -m llama.gguf

A GGUF file holds weights and the name of an architecture, and llama.cpp runs it with its own implementation of that architecture, so a model can only be written as one whose computation is the source’s. --gguf finds which, of llama and gpt2 (or --arch), and which field is which of its tensors, by computing both on random weights in f64: the forward function as the source writes it, and the architecture as llama.cpp computes it. Fields are paired with tensors in the order the forward function first reads them, then in the other orders of fields of one shape, and for llama, with each common rotary base, and with the query and key rows reordered as convert_hf_to_gguf.py reorders Hugging Face’s, for rotary embeddings that turn the two halves of a head rather than neighbouring features. The first pairing whose logits agree to 1e-8 of their size is written. If none does, nothing is written, and the error says why for each architecture: a different norm, a different number of tensors, or the closest pairing’s difference, as for GELU computed exactly where llama.cpp’s gpt2 uses tanh. Every norm’s eps must be one number written in the calls, which the file holds once.

• llama: rms_norm, rotary position embeddings on neighbouring features, causal attention with as many key and value heads as the key projection has rows for, and a SwiGLU feed-forward layer. examples/llama.ss is one.
• gpt2: layer_norm, learned position embeddings, a fused query, key and value projection with biases, and GELU with tanh. examples/gpt2.ss is one.

Matrices are written in --dtype, and 1-D tensors in f32. --tokenizer DIR writes a byte-level BPE tokenizer, such as GPT-2’s, from tokenizer.json; without one, the file’s vocabulary names each token <|N|>, which llama.cpp needs, and prompts are token ids. llama.cpp’s own kernels compute in lower precision in places, such as half-precision attention on a GPU and a GELU table on the CPU, so its log-probabilities differ from the source’s by up to about 1e-2 on a GPU and 1e-3 on the CPU, with its cache in f32.

Serving exports from Python

shapescript.export_hf(path, output, *, fn="forward", weights=None, seed=0, dtype="float32", tokenizer=None, max_len=None, verify=True, vllm=False) writes the folder and returns the config.json it wrote; weights may also be a model. shapescript.export_gguf(path, output, *, fn="forward", arch=None, weights=None, …) writes the GGUF file and returns its architecture, or raises ValueError saying why the model isn’t one. A file with errors raises shapescript.ShapeScriptError.

Notebooks

%load_ext shapescript adds the %%shapescript cell magic to Jupyter and VS Code notebooks. A cell is a file that may also have statements outside any fn.

Each cell is checked together with the cells that have run without errors, in the order they last ran: declarations are hoisted as in a module, a name declared in several cells has the latest cell’s declaration, and statements outside any fn are checked in order, in one scope. A for or if outside any fn follows the block rules within its cell; a later cell sees the values first bound in the block as Python values. Python variables of the notebook are visible as Python values where nothing else has the name. A cell with errors shows them as its output and runs nothing; a cell that ends with a tensor shows its shape diagram.

What a cell made can’t be used once it doesn’t match the cell’s code any more, until the cell runs again (E0701): what a cell defines when the cell no longer checks, a value whose type isn’t the one it was made with (as after its dims changed, including a model whose fields’ shapes changed), a declaration that Python code has assigned or deleted, and functions, domains and models that use any of those.

Running programs

shapescript run [--python PY] [--] program [args…] runs program with PY -m shapescript program args… (see Debugging): a Python file, which runs as python program.py would but can import .ss files without shapescript.install(), or a .ss file, which is imported and whose fn main() is called. sys.argv is [program, args…]. The Python is --python, else $SHAPESCRIPT_PYTHON, else python3 (python on Windows), as for shapescript export. The command’s exit code is the program’s.

Running on Modal

With --modal GPU, the program runs in a container on Modal, with that GPU, instead of here; the guide is docs/modal.md. Nothing in the program changes. The command needs the modal package in that Python (pip install modal) and a Modal account (modal setup); without the package it says so and exits 2.

GPU is Modal’s name for a GPU type, in any case: T4, L4, A10G, A100, A100-80GB, L40S, H100, H200, B200 and others, with :N for N of them, 1 to 8, as in H100:3. Inside the container, devices cuda[0] (and cuda[0 1 2] with three GPUs) is available. A program that declares devices cpu or cpu[…] runs on the CPU, and one that declares devices mps runs on the CPU too, since Modal’s machines are Linux; both print a warning that the GPU is unused.

The container has the local Python version, torch, numpy and einops at the local versions (Linux builds, with CUDA), and this version of ShapeScript: from PyPI when it is published there, else built from the source checkout the local one was built from. --pip REQ installs another package, or another version of those three, and --shapescript says which ShapeScript: a Linux x86_64 wheel, a source checkout, or a pip requirement. The first run builds this image, which takes minutes; Modal keeps it for later runs.

The program’s folder, or --project DIR, is uploaded to the same absolute path in the container, without .git, .venv, venv, __pycache__, *.pyc, target, node_modules and tools’ caches, nor what matches an --exclude GLOB (a pattern without / matches any file or folder name, one with / a path in the folder). More than --max-upload MB (500) of it is an error that names its largest entries. The program runs in the same working directory as the command, so relative paths, __file__ and tracebacks are the local ones: a ShapeScript error raised in the container names the local .ss file and line.

What the program prints is streamed back as it prints it, its stdout to stdout and its stderr to stderr; Modal’s progress and the image’s build log go to stderr. When the program ends, the files it created or changed in its folder, or in the working directory when that contains its folder, come back to the same paths here; files it deleted are kept here. The exit code is the program’s.

Ctrl-C stops the program and its container, and the command exits 130; the files it wrote don’t come back. A .ss program is compiled here before anything starts, so its errors cost nothing, and so is a devices line that asks for more CUDA devices than the GPUs: devices cuda[0 1 2] with --modal T4 is an error that suggests --modal T4:3.

Completion and signature help

The language server completes names everywhere a name can go, the way gopls does for Go: as you type, with the candidates whose types fit the place first.

What’s offered where.

• In an expression: the values in scope where the cursor is (the function’s parameters, the values its statements bound before the cursor, a for loop’s variable inside its body, and after an if a value both branches bind), then the file’s dims, functions, models, records, imports and import python modules, then every built-in, the prelude’s included, and Train(seed=…), Eval, true, false and inf. Each has its kind, its type or signature as the detail (y: [b n D], dim Dh = D / H = 4, fn softmax(x: [... n] Logit over n, over: Axis[n]) -> [... n] Prob over n), and the hover’s Markdown, ## docs included, as its documentation, sent when the editor asks for it (completionItem/resolve). A function inserts a call with a tab stop for each parameter without a default, keyword ones by name: softmax(${1:x}, over=${2:n}); to an editor without snippets, just its name.
• In a call’s arguments: after ( or ,, the keyword arguments not given yet (over=, weight=, eps=, a stub’s parameter names), which open the list again once chosen; after over=, exactly the axes of the operand (the piped value, or the first argument); for a Mode parameter, Train(seed=…) and Eval; for a mut parameter, mut.
• After |>: the functions whose first parameter takes the piped value, as stages without it: softmax(over=${1|D,n,b|}) with the value’s axes to choose from, last first, or gelu alone. A function whose first parameter can’t take the value isn’t offered.
• In a pattern of rearrange, reduce, repeat, einsum or split: before ->, the operand’s axes, all of them as one item first, then the next one; after it, the names the inputs have. After reduce(x, b n D -> b D, , the reductions.
• In a type: where a type starts (after x: , -> , as or to ), a shape, the file’s models, records, logs records, shape aliases and domains, and Mode, int, float, bool, str, Grad[M], AdamW[M], SGD[M], Dataset[E], Schedule, LogSink, Axis[n], Axes[n] and Size[k]; after as or to, the value’s own type first. Inside […], the axes the signature names, the dims and shape aliases, and ...; after the shape, the domains and dtypes; after over, the shape’s axes; inside Index[, Axis[, Axes[ and Size[, axes and dims, and inside Grad[, AdamW[ and SGD[, models.
• After a .: a model value’s fields and functions (completion after m.); a batch’s fields with their types (field img: [b D]), a logs record’s metrics, a gradient’s params; an imported file’s declarations (attention.); init after a model’s name (Mlp.init(seed=…)); and a Python module’s names (torch., torch.nn.functional.), with their signatures and docstrings.
• After import : the .ss files of the file’s folder, the workspace’s folders and the search path, as dotted paths a part at a time (models. then mlp), the standard library and installed packages with .ss files (Packages), and python. After import python : the Python modules on the interpreter’s path, and a package’s submodules after torch..
• Keywords where they go: at the top level, declarations as snippets (fn name(x: T) -> T: with its block, model Name: with a param line, data, logs, dim, domain, extern fn, export fn, import, import python, devices); at the start of a statement, for i in 0..n:, if cond:, else: and return, with the names; after a value and a space, as, to, and, or, for, if and else, and in after for x .
• Names of other files: once a letter is typed, the declarations of the workspace’s other .ss files, the standard library’s and installed packages’ that the file doesn’t import yet, qualified (metrics.accuracy), with an edit that adds import metrics after the file’s imports, or before its first declaration. They come after the names in scope, marked metrics · adds its import. A file that imports this one isn’t offered, which would be a cycle.

Ranking. The place’s expected type comes from the parameter of the call the argument is for (the overload its first argument’s model picks), the function’s return type after return, the other operand of an operator (one whose axes are the same or fewer broadcasts; for @, a matrix whose first axis is the left’s last), the parameter after name=, and the piped value for a stage’s first parameter. A candidate fits exactly if it has the same shape with the same dims (a [b n D] for x: [b n D]), the same model, record or mode; it fits if its rank and dims agree with a generic or ... shape; it doesn’t if the call would be an error. Exact fits come first, and the first is preselected; then those that fit; then the rest, values in scope before the file’s declarations, built-ins, keywords and other files’ names; and last those that don’t fit. A function’s fit is its return type’s. Within each, the editor’s own matching orders them, and the server’s fuzzy matching decides what’s in the list: the letters typed must appear in the name in order, ignoring case, and matches at the starts of words (lnorm for layer_norm) and runs of letters score higher. A list is marked incomplete where typing more brings other items (other files’ names, a Python module still being listed, more than 300 items), so the editor asks again.

Edits. An item replaces the word typed so far, from its start to the cursor, or the whole word around the cursor if the editor supports both; after ., = or \, the word after it.

While the file is broken. A line being typed rarely parses. The names in scope come from the text with the cursor’s line left out, parsed, so the statements before it count, and their types from the last analysis of the file without syntax errors, matched by name and place; the type of an expression on the line itself, as pool(y) + 1 in mean(pool(y) + 1, over=, comes from checking a copy of the file cut after it, with its brackets closed. On a file of 2,400 lines, a completion takes about 3 ms in a release build and 12 ms in a debug build, 22 ms over the protocol.

Python. The modules on the path come from the interpreter that go to definition uses, listed with pkgutil without importing anything. A module’s names come from its source, read with ast, when it’s one of yours (under the file’s folder or the workspace’s), so completing a name never runs your code; others, such as torch and the standard library, are imported and asked with dir() and inspect, once per interpreter for the session.

Signature help (textDocument/signatureHelp, after (, , and =) shows the signature of the call being typed, with the parameter being typed highlighted and its type: the positional one for the argument, the one named before =, or once the positional arguments are given, the keyword one still missing, as over. In a pipe stage, the first parameter is the piped value. Functions of the file, overloads (all of them, with the one the first argument’s model picks active), generic functions, imported files’ functions, built-ins with their documentation, extern fns, Mlp.init and Python functions (their inspect or ast signature) all have it.

Trigger characters. The list opens by itself after .; after [ in a type; after ( and , where the argument’s type is known; after = that follows a keyword argument’s name; after a space that follows |>, import or import python; and after \ and { in math in comments. Anywhere else these characters open nothing.

Debugging

ShapeScript is debugged with Python’s debugger, debugpy, which VS Code’s Python Debugger extension and Jupyter use. The code the import hook, python -m shapescript and the cell magic run reports the .ss file (or the cell) and its lines, so a breakpoint in a .ss file stops there, the call stack mixes .ss and Python frames, and stepping goes from one language to the other: step into a ShapeScript line that calls Python and you’re in the Python function; step out and you’re back on the line; step into a call from Python and you’re on the fn line of the ShapeScript function.

Every line is a place to stop. A pipe written over several lines stops on the line before its stages and on each stage. The fn line is where the argument checks run.

When Python imports a .ss file that has errors, through the import hook or the module generated beside it, the ShapeScriptError it raises, which shows every diagnostic, is raised from the file’s first error: the last entry of its traceback is that line of the .ss file, as File "model.ss", line 5, in <module>, with the error’s span underlined from Python 3.11 on. So a debugger that stops on uncaught exceptions, as VS Code’s does by default, stops on the broken line of the .ss file and shows the error there, not in the Python that imported it.

pip install shapescript installs a debugpy extension (pydevd_plugins/extensions/pydevd_plugin_shapescript.py), which debugpy loads whatever starts it. It:

• hides the names the compiler makes up, which all start with _ss_ (temporaries, sizes, the device, the runtime module);
• shows a tensor in a .ss frame with its sizes, named after the axes of its declared type, before its values: x = [b=2 D=4] float32 tensor([[1., …, and gives each parameter and value its ShapeScript type, as hover does ([b D] Prob over D), which VS Code shows when you point at the variable, and so each field of a model it holds (m.w);
• does the same for functions that share a name (overloaded functions), whose Python names are _ss_forward__Mlp and the like, and for each instance of a generic function, _ss_fit__Mlp__Digit, whose values have the types of that instance; an instance’s code is shown at the call it was made for, as its lines are. Imports of other .ss files are found as the import hook finds them;
• steps over the shapescript package, the runtime that generated code calls, even with "justMyCode": false or an editable install. A launch configuration’s rules come first, so {"module": "shapescript", "include": true} steps into it.

python -m shapescript program.py [args…] runs a Python program with the import hook installed, so it needs no shapescript.install(); python -m shapescript model.ss [args…] imports the file and calls its fn main(), which takes no parameters; python -m shapescript -m module runs a module. VS Code’s shapescript launch configurations run these under debugpy (see the extension’s README); any other debugpy configuration works too, such as the Python extension’s “Python File”, if the program installs the hook itself.

In a notebook, the %%shapescript cell magic notices when Jupyter’s debugger (ipykernel’s, which VS Code and JupyterLab use) has written the cell to a file to set breakpoints in it, and its code then reports that file and the cell’s lines, counting the %%shapescript line, so breakpoints in the cell stop. Types come from the cell alone, so values whose types use what earlier cells declared show only their sizes.

Not debuggable yet: the worker processes of several devices, where each device’s work runs (debugpy debugs only the controller’s process, so breakpoints in fns that run on the workers don’t stop); and Python written by shapescript build, which is plain Python whose lines are its own.

The tensor view

In VS Code, a stopped program’s tensors are drawn in the Tensor view, a section of the Run and Debug view beside Variables, Watch and Call Stack, which shows while debugging. View Tensor adds an expression to it and shows it: from the context menu of a variable in the Variables or Watch view, from the editor’s context menu on a name (or a selection) while stopped, or from the Command Palette or the + in the view’s title bar, which ask for an expression. The view keeps its expressions as Watch does, one tab each, and draws the selected one; its expression box changes it. The expression is evaluated in the selected frame, as a watch is, so it can be any Python expression of the frame’s names, such as m.w or att[0], and a NumPy array works too. When the program stops again, or another frame is selected in the call stack, the view shows the expression there; while the program runs, it keeps what it drew, marked as the values of the last stop. In the side bar it’s one column, the heatmap, then the checks and statistics; made wider, or moved to the panel, it puts them side by side.

The tensor is drawn as a heatmap of two of its axes, zoomed with the mouse wheel and panned by dragging, from the whole tensor down to single elements, which show their values as numbers once the cells are big enough. Pointing at a cell shows where it is, by axis name (att[b=0, H=2, n=20, m=41]), and its exact value, written with the fewest digits that are the same number in the tensor’s dtype; a cell that covers several elements shows their mean, least and greatest. Each axis has a role, chosen in the view: the heatmap’s rows or columns, small multiples (a heatmap for each of its indices), color channels (for an axis of 3 or 4), a slider, or a reduction over it (mean, max, min, sum, largest |x|, norm or std). A group axis whose parts’ sizes are known, such as (gh gw), can be split into its parts, which then have roles of their own. A tensor of one axis is drawn as a line, with the least and greatest values under each pixel as a band, and a scalar as its value.

The tensor’s ShapeScript type, from the checker, as hover shows it, picks the layout and the colors. Values of other Python code, which have no type, are drawn from their dtype alone, with their axes numbered.

• An axis made of two named sizes followed by one more axis, as the tokens of a grid of patches [b (gh gw) D], is split, and its parts are the rows and columns, each cell the token’s norm over D.
• Otherwise, where the checker knows a matrix is a background outside a band of diagonals, as a causal mask causal[n m] is false above its diagonal, and what masked_fill and softmax make of one is -inf or 0 there (its structure, which hover draws), its rows and columns are the heatmap’s.
• Otherwise Logit, Prob or LogProb over m puts m on the columns and the axis before it on the rows.
• Otherwise an axis of size 3 before the last two is the color channels of an image whose values are Unit or u8, and a slider for other values; the last two axes are the rows and columns.
• Otherwise the last two axes are the rows and columns.

Then the nearest remaining axis before the rows, if it has 2 to 16 elements, is small multiples, as H is in [b H n m], and the rest are sliders at 0. Colors: a mask (Mask, Binary) has two; Index[V] with at most 8 values has the categorical palette’s colors; values that can’t be negative (Prob, Unit, NonNeg, LogProb) have a sequential scale from the color nearest the background, so small values recede; and others have a diverging scale around 0, blue below and red above. The range is the domain’s when it’s bounded, as [0, 1] for Prob; otherwise the values’ 0.5th to 99.5th percentiles while each cell is an element, and the shown cells’ range when cells are reductions or cover several elements. The histogram beside the heatmap is its legend: dragging its ends sets the range, and double-clicking resets it. A cell that covers several elements shows, by default, the greatest of values that can’t be negative, the value of largest |x| (with its sign) of others, and the mean of masks and indices, so that a single outlier shows at any zoom; the view’s Cells menu chooses the mean, least or greatest instead.

The checks list what the type promises and whether the values keep it, with Show, which goes to the first or worst element that doesn’t:

• no NaN or ±∞ (-inf where a mask’s background is -inf, or in Additive values, isn’t counted);
• every value is in the domain: in its range, as [0, 1] for Unit or Prob, at most 0 for LogProb, 0 or 1 for Binary, 0 or -inf for Additive, and whole numbers in [0, V) for Index[V];
• each row of Prob over m sums to 1 along m (each row of LogProb’s exponentials does), to within 1e-3 or twice the dtype’s epsilon, whichever is larger;
• the background its structure promises is there: after a causal mask, att[b, n, m] is 0 wherever m > n.

On the heatmap, NaN and ±∞ are hatched in red, values outside the domain, and those that aren’t the background their structure promises, in orange, and that background is hatched in gray, with its edge dashed; rows of a Prob that don’t sum to 1 are marked at the heatmap’s left edge. Where cells are smaller than a few pixels, a ring marks each flagged one, so one NaN in a billion values shows. Beside the checks are the values’ statistics, and, for floating-point dtypes, how far the largest |x| is from the dtype’s largest number (its headroom, in bits) and how many values are subnormal.

The tensor never leaves the device it’s on. View Tensor evaluates shapescript._tensorview.open(...) in the frame, which keeps a reference to the value and starts a server in the program, on 127.0.0.1 with a random token; the view asks it for a summary and for tiles of 256 × 256 cells, which it computes where the tensor is, reading it in blocks whose temporaries take at most an eighth of the GPU’s free memory, and sends only those. A cell covers 2ᵏ × 2ʲ elements, so a few tiles cover the whole tensor when zoomed out, and the cells keep each element’s NaN and ±∞ and whether it’s in the domain. On a 5090, a summary of 8 GiB of bf16 takes under a second, and a zoomed-out tile of all of it about 160 ms. The server’s threads are marked as the debugger’s own, so debugpy neither traces nor suspends them, and on CUDA they read on a stream of their own after the work the stopped thread had queued, so the program’s streams are untouched. A value that’s being recorded into a CUDA graph has no values yet, and isn’t drawn. Random samples, for the percentiles, use a generator of their own, so the program’s random numbers don’t change. When the program goes on, the view releases the value; when the debugger goes away, the server releases those it read.

Not yet: values on another machine than VS Code’s extension (with Remote SSH and containers, the extension runs where the program does), complex, sparse and quantized tensors, values inside a map (which PyTorch’s vmap sees as batched tensors), and the worker processes of several devices.

Not yet implemented

Training on one device is implemented: control flow, host scalars, models, mut parameters, gradients, modes and dropout, data records and datasets, logs, sinks and print, optimizers, gradient transforms and schedules, checkpoints, the standard library’s fit and the devices declaration. So is data parallelism on several devices, with a controller and a worker per device, with Python calls checked and extern fns run on each device’s examples, and the options shard params and split H; all of it has been run on worker processes on the CPU (devices cpu[0 1 2]) and, with NCCL, on 3×T4 and 3×A100 GPUs on Modal, where it matches one GPU’s results. Placement and communication show on hover. Steps are recorded as CUDA graphs, with new fixed shapes recorded ahead of time, and on each worker of several GPUs. A step on several GPUs is as fast as PyTorch’s DDP when the dataset fits on the GPUs, but not yet with large batches loaded from the CPU (see Speed on several devices). What isn’t done yet, beyond milestone M10 in IMPLEMENTATION_PLAN.md: distributed lowerings of reductions over the batch other than sums and means; several devices in notebooks; and dropout masks that match one device’s. Also not yet: x to cpu and x to device, and random operations other than dropout. The mask constructors tril, triu, eye and band are planned too. The only built-in optimizers are SGD and AdamW, the only built-in gradient transform is clip_norm, and the only schedules are cosine and warmup.

Diagnostics

Every diagnostic has a stable code. The letter is its severity: E for errors, which stop a build, W for warnings and H for hints. The first two digits are its area: 01 lexing and layout, 02 parsing, 03 names and declarations, 04 shapes, 05 value domains and dtypes, 06 math in comments, 07 notebooks, 08 control flow and host scalars, 09 models, 10 devices, 11 mut and gradients, 12 modes and random operations, 13 data records and datasets, 14 logs and sinks, 15 optimizers, gradient transforms and schedules, 16 checkpoints, 17 export, 18 several devices, 19 Python calls and options on several devices, 20 fixed shapes and padding, 22 grouped sizes and shape aliases, 23 modules, overloaded functions and generic functions, 24 iterator chains: what they iterate and their independence (comprehensions before CTL-3), 25 the Python modules generated beside .ss files, 29 installed packages, 31 recording steps as CUDA graphs, 32 iterator chains: their syntax and stages, 33 exporting models for Hugging Face and serving engines. Codes are never reused.

Each section explains a code, shows a short program that produces it, and links to the rules it enforces. shapescript explain <code> prints a section in the terminal.

01: Lexing and layout

E0101

Unexpected character. A character that can’t start any token, such as $, % or ; outside a comment or string. A run of them is reported once.

dim D = 768 $

See Tokens.

E0102

Unterminated string. A string has no closing " before the end of its line. Strings can’t span lines.

extern fn erfc(x: Real) -> NonNeg = "torch.special.erfc

See Tokens.

E0103

Tab used for indentation. Indentation must be spaces. Tabs elsewhere on a line are fine.

fn f(x: [n]) -> [n]:
	return x

See Lines and indentation.

E0104

Indentation doesn’t match an enclosing block. A line is indented less than the line before it, but not as little as any block it could be closing. The help says which indentations would work.

fn f() -> []:
        a = 1.0
    return a

See Lines and indentation.

E0105

Unmatched closing bracket. A ) or ] with no opening bracket before it.

fn f(x: [n]) -> [n] = add(x, x))

See Lines and indentation.

E0106

Unclosed bracket. A ( or [ is never closed. Since line breaks don’t count inside brackets, the rest of the file is read as part of the bracket, which usually causes other errors too.

fn f(x: [n]) -> [n] = add(x, x

See Lines and indentation.

02: Parsing

E0201

Expected a different token. The parser found a token where the grammar doesn’t allow one: a missing operand, a missing : or = after a signature, a declaration inside a function, a return outside a function in a notebook cell, and so on. The message says what was expected.

dim H = * 2

See Grammar.

E0202

Unexpected indentation. A line is indented more than the line before it, where no block starts. Only the line after a : that ends a fn, for, if, else or model line, or after a devices line, starts an indented block; to continue an expression on the next line, put it in brackets or start the line with |>.

fn f(x: [n]) -> [n]:
    return x
        extra

See Lines and indentation.

E0203

Pattern outside a pattern operation. A pattern (... -> ...) can only be an argument to rearrange, reduce, repeat, einsum or split.

fn f(x: [n D]) = softmax(x, n D -> D n)

See Expressions and Pattern operations.

E0204

Invalid pipe stage. A pipe stage must be a function name, a call, or to T. Since a stage is parsed like an operand of or, x |> f + 1 makes f + 1 the stage. Wrap the pipeline in parentheses to use its result: (x |> f) + 1, (x |> f) as T.

fn f(x: [n]) = x |> relu + 1

See Pipes and Operator precedence.

E0205

Misused _. _ can only be an argument of a call that is a pipe stage, and only once per stage.

fn f(x: [n]) = add(_, x)

See Pipes.

E0206

Chained comparison. Comparisons don’t chain: a < b < c isn’t a < b and b < c. Write that instead.

fn f(x: [n]) = 0.0 < x < 1.0

See Operator precedence.

E0207

Expected a declaration. At the top level of a file there can only be devices, dim, domain, fn, extern fn, import python and model. Statements belong inside a function, except in a notebook cell.

x = 1

See Files and items.

E0208

Integer too large. An integer literal is larger than 18446744073709551615 (2^64 − 1).

dim D = 99999999999999999999

See Tokens.

W0201

Doc comment not attached to a declaration. A ## comment documents the declaration on the line directly below it. This one has none: it is followed by a blank line or a statement, or it is at the end of a line of code. Move it, or use # for an ordinary comment.

## The model width.

dim D = 768

See Comments and doc comments.

03: Names and declarations

E0301

Unknown name. A name that nothing declares: a value, function, dim, domain or axis. When a name with a similar spelling exists, the help suggests it. In a return type or an as, a size name that isn’t in scope is reported as an unknown size; write ?k for a size the function chooses or that is only known at run time.

fn f(x: [b D]) -> [b D] = relu(x) |> gelo

See Names and scoping and Axis names.

E0302

Name defined twice. Two module-level declarations, two parameters of one function, or two fields of one model have the same name.

fn f(x: [n]) -> [n] = x
fn f(x: [n]) -> [n] = x

See Module scope.

E0303

Name used as the wrong kind of thing. A dim or value called as a function, a function or domain used as a value, a value used as a size, a value named like an axis or a dim, a dim used as a domain, and so on. The message says what the name is.

dim D = 768

fn f(x: [n D]) = D(x)

See Function scope and Name clashes.

E0304

Dim defined in terms of itself. Dims may refer to each other in any order, but not in a cycle.

dim A = B + 1
dim B = A * 2

See Dims.

E0305

Dim isn’t a positive whole number. A dim’s value must be at least 1, and a division in it must be exact. The diagnostic points at both dims when a division fails.

dim D = 768
dim H = 10
dim Dh = D / H

See Dims and Division.

E0306

Invalid function signature. A parameter with neither a type nor a default; ?k in a parameter type; in a fn, a default that isn’t a number, that uses the function’s parameters or axis names, or that is on a parameter with a shape or a model’s type; a default that isn’t a number on an untyped parameter of an extern fn; ..., Axis[..], Axes[..] or Size[..] in a fn, or ... twice in one type or inside a group; a call to a function without a return type from inside its own body.

fn f(x) -> [n] = x

See Parameters, Unknown sizes and Extern functions and the prelude.

E0307

Missing return. The signature declares a return type, but the body doesn’t return a value on every path, or has a bare return. A for loop may run no times, and an if without an else may not run, so a return in them doesn’t count; put one after the loop, or in both blocks of an if.

fn f(x: [n]) -> [n]:
    y = x

See Function declarations and Control flow.

E0308

Invalid extern fn. An extern fn without a return type; a string that isn’t a Python function’s dotted name (it needs a module and a function, each part a Python identifier and not a keyword); an Axis[..] or Axes[..] parameter with a default or without a first parameter with a shape; a Size[..] default that isn’t a number or an earlier size; a return type with ... when no parameter has it.

extern fn erfc(x: Real) = "torch.special.erfc"

See Extern functions and the prelude.

04: Shapes

E0401

Sizes don’t match. Two sizes that must be equal can’t be: an argument’s axis and its parameter’s, a returned value’s and the return type’s, two axes that broadcast together, the axes @ contracts, a pattern’s name or number and the axis it binds, the parts of a pattern group and their axis. The notes say which axis and why, as in “?k is not known to equal 196”. A generic size or an unknown size is only equal to itself, so n doesn’t match 196 and ?k matches nothing else; use as to claim a size the checker can’t prove.

fn f(x: [n 3], w: [4 5]) = x @ w

See Unification and the size solver, Generic calls and Broadcasting by name.

E0402

Size isn’t divisible. A size divides by something that doesn’t divide it exactly: in a type or dim (D / 5), when solving a pattern group or a call ((gh 16) against 225), or in an extern fn’s rounded division with a size that isn’t known.

dim D = 768

fn f(x: [n D]) -> [n D / 5] = x

See Division, Solving and Rounded division.

E0403

Wrong number of axes. A tensor has a different number of axes than its parameter, the return type, the pattern side describing it, or the tensors it is joined with; @ got a tensor without axes; or there are more indices than axes.

fn f(x: [n D]) = rearrange(x, n -> n)

See Generic calls, Pattern operations and Indexing.

E0404

Axes can’t be lined up by name. Neither operand’s names are all axes of the other, so the axes line up from the right, and two named axes meet with different names. The same holds for the tensors of cat, stack, gather and scatter, which line up by position, and for a masked_fill mask with axes the tensor doesn’t have. If the axes really are the same, give them the same name with rearrange or as.

fn f(x: [n D], y: [k D]) = x + y

See The order rule.

E0405

Tensor has no such axis. over= names an axis the tensor doesn’t have, or a pattern’s output names an axis that isn’t in its input. The note lists the tensor’s named axes.

fn f(x: [n D]) = softmax(x, over=q)

See Built-in operations and Pattern operations.

E0406

Wrong arguments. Too many or too few arguments, an unknown keyword, an argument given twice, a missing or invalid over=, an unknown reduction, a bad index or slice bound, a keyword size for a pattern name that already has one, an Axis[..] argument that isn’t a name, a Size[..] argument that isn’t known statically, and so on.

fn f(x: [n D]) = softmax(x)

See Matching arguments and Built-in operations.

E0407

Pattern doesn’t fit the operation. An input name missing from the output of rearrange, repeat or split; a number on the right of -> that isn’t on the left; a name twice on one side; a nested group; a group or number in einsum; a pattern operation without a pattern or with two; split without exactly one number that gives the count.

fn f(x: [b n D]) = rearrange(x, b n D -> b D)

See Pattern operations.

E0408

Axes in a different order. A pattern names an axis by a name the operand has at another position. Patterns bind by position, so this is probably a transposition; reorder the axes with rearrange first.

fn f(x: [b n D]) = rearrange(x, b D n -> b n D)

See Pattern operations.

E0409

Size can’t be inferred. The size solver is left with unknowns that no equation determines: two new names in one pattern group, a new axis in repeat without a size, or a call whose argument doesn’t determine a grouped parameter’s parts. Give all but one of them a size, as in h=8.

dim D = 768

fn f(x: [b n D]) = rearrange(x, b n (h e) -> b n h e)

See Solver failures and Pattern operations.

E0410

as that can never succeed. The asserted shape has a different number of axes or a different known size, the asserted domain contradicts what’s known about the values, or the asserted dtype is of another kind. An as claims what the checker can’t prove, but not what it can disprove.

dim D = 768

fn f(x: [n D]) = x as [n 5]

See Assertions with as.

E0411

Wrong number of results. split and topk give several results, which must be bound to as many names, and a single value can’t be bound to several names.

fn f(x: [n 12]) -> [n 4]:
    a, b = split(x, n (3 e) -> n e)
    return a

See Statements and Pattern operations.

E0412

Value isn’t a tensor. A string, or the several results of split before they are named, where a tensor is needed.

fn f(x: [n]) = x + "s"

See Operators.

E0413

Python value with an unknown shape. A Python value is used where its shape or type must be known: returned as a typed result, passed as a tensor argument, given to a pattern operation or an operation along a named axis, or converted with to and a domain. Assert its type with as first.

import python torch

fn f() -> [8] = torch.ones(8)

See Python interop.

E0414

Size too large. A size, or a step in working one out, doesn’t fit in 128 bits.

dim Big = 4294967296 * 4294967296 * 4294967296 * 4294967296 * 4294967296

See Size polynomials.

E0415

More elements than the axis has. topk asks for more values than its axis has, or a negative number.

dim C = 10

fn f(x: [b C]) -> [b C]:
    v, i = topk(x, 11, over=C)
    return x

See Built-in operations.

05: Value domains and dtypes

E0501

Value isn’t in the expected domain. A value doesn’t fit where it’s used: a range not contained in the expected one, a Prob where logits are expected, a Mask where numbers are expected or the other way round, and so on. The help suggests to T when a conversion exists and as T otherwise.

fn f(x: [n]) -> [n] Unit = x * 2

See Widening, Semantic domains and Discrete domains.

E0502

Wrong dtype. An argument or returned value is stored in a different dtype than expected, or the operands of @, matmul, einsum or scatter have different dtypes. Convert with to, as in to f32.

fn f(x: [n D] bf16, w: [D D]) = x @ w

See Dtypes.

E0503

Invalid domain. Logit, Prob or LogProb without over, or over an axis the type doesn’t have, or in a type without a shape where none can be found; over or a size on a domain that doesn’t take one; Index without a size; a domain and dtype that don’t go together, such as Mask f32; a domain declaration with a reserved name, ends that aren’t numbers, or no values.

fn f(x: [b C] Prob) = x

See Built-in domains and User domains.

E0504

to conversion that doesn’t exist. to only makes the conversions in its table, and the note lists them. Converting to the domain a value already has, converting semantic or discrete domains by their values’ range, and converting to bool are errors too.

fn f(x: [n]) = softmax(x, over=n) to Signed

See The to table and Converting dtypes.

E0505

Arithmetic on a mask. + - * / @ and unary - take numbers, and a Mask holds booleans. Convert it: to Additive to add it to scores, to Binary for 0 and 1; use ! to negate it and and or or to combine masks.

fn f(x: [n m], keep: [n m] Mask) = x + keep

See Discrete domains.

E0506

Function applied outside its domain. log, log1p, sqrt, rsqrt or a division gets values that are past where the function is defined, or a divisor that is exactly 0.

fn f(s: [n] Signed) = log(s)

See Function edges.

E0507

Indices for an axis of a different size. Index[V] values index an axis whose size isn’t V: in emb[ids], a call, gather, cross_entropy targets, or a return type.

dim V = 100
dim V2 = 50

fn f(emb: [V 8], ids: [t] Index[V2]) = emb[ids]

See Discrete domains and Indexing.

W0501

Value can be exactly where a function is undefined. The range of the argument of log, log1p or rsqrt, or of a divisor, reaches the function’s edge: log of values that can be 0, or a division by values that can be 0. The help suggests a stable alternative where there is one, such as log_sigmoid(x) for log(sigmoid(x)), or as Pos if the value is never 0.

fn f(x: [n]) = log(sigmoid(x))

See Function edges.

06: Math in comments

H0601

LaTeX outside the supported subset. Math in a comment uses a command or environment that isn’t in the subset. It is shown as written. A misspelling of a supported command gets a suggestion.

## Normalizes with $\sofmax(x)$.
fn f(x: [n]) -> [n] = x

See Math in comments.

H0602

Malformed math in a comment. Math in a comment can’t be parsed: an unclosed {, a missing argument, & outside an environment, an unmatched \left or \end, two superscripts in a row. It is shown as written.

# Scaled by $\frac{1}{\sqrt{D}$.
fn f(x: [n]) -> [n] = x

See Math in comments.

07: Notebooks

E0701

Something from a notebook cell that has to run again. A cell uses a value, function or domain from an earlier cell that no longer matches that cell’s code: the cell doesn’t check any more, a value’s type has changed since it was made (as after a dim it uses changed), Python code assigned or deleted it, or it uses something that is out of date. Run the cell named in the message again.

# In [1]
dim D = 4
# In [2]
x = torch.ones(2, 4) as [2 D]
# In [1], changed and run again
dim D = 8
# In [3]
y = x * 2            # error: `x` is out of date; run [2] again to use this

See Notebooks.

08: Control flow and host scalars

E0801

Value assigned in a loop or if changes its type. A loop or if block assigns a value from outside it, or a loop’s variable has the name of one, and the new value has a different shape, different axis names or a different dtype, is another model, or is another kind of host scalar. The code after the block, and the next iteration of a loop, must work either way, so only the value’s domain may change. Give the new value a new name.

dim D = 64

fn f(x: [b D], w: [D 2*D], n: int) -> [b D]:
    for i in 0..n:
        x = x @ w
    return x

See Blocks and scope.

E0802

if on something other than a bool. An if condition must be a bool, a [] Mask or a Python value. A mask with axes isn’t true or false: reduce it with any or all. A number must be compared with something, as in n > 0.

dim D = 64

fn f(x: [b D]) -> [b D]:
    if x > 0:
        x = x * 2
    return x

See Conditions.

E0803

Range whose ends aren’t ints. The ends of a..b in a for must be ints or Python values. A float isn’t one, and neither is a tensor, even a [] i64 one, since a range counts on the host.

dim D = 64

fn f(x: [b D], n: float) -> [b D]:
    for i in 0..n:
        x = x * 2
    return x

See Loops and Host scalars.

E0804

for over something that can’t be looped over. A for loops over a range, as in for i in 0..n:, over a dataset’s batches, as in for batch in train |> batches(64):, or over a Python value. A tensor can’t be looped over: index it in a loop over a range instead. A dataset is looped over a batch at a time.

dim D = 64

fn f(x: [n D]) -> [n D]:
    for row in x:
        x = x * 2
    return x

See Loops and Looping over batches.

E0805

Tensor where an int, float or bool is expected. A host scalar parameter, a host scalar return type, as int, as float or as bool, or the seed of init got a tensor or a model. Host scalars are Python values on the host; a tensor, even one of shape [], lives on the device, and reading it would wait for the device. There is no conversion from a tensor to a host scalar.

dim D = 64

fn scale(x: [b D], times: int) -> [b D]:
    for i in 0..times:
        x = x * 2
    return x

fn f(x: [b D]) -> [b D] = scale(x, sum(x) to i64)

See Host scalars.

09: Models

E0901

Model has no such field. m.name reads a field that the model doesn’t declare. The help suggests a field with a similar name.

dim D = 64

model Mlp:
    param w1: [D D] = normal(std=0.02)

fn f(m: Mlp, x: [b D]) -> [b D] = x @ m.w2

See Using models.

E0902

Invalid model field. A field’s type has ?k, has no shape, is a host scalar, a model, optimizer state or a Schedule; a param holds something other than floats or has a domain; a buffer has a domain other than a range, Index[V] or Mask; or the field’s name is init, a Python keyword, starts with _, or is an attribute of torch.nn.Module.

dim D = 64

model Counter:
    param count: [] i64 = zeros

See Fields.

E0903

Invalid field initializer. A field has no initializer, or one that isn’t zeros, ones, full, normal or uniform; its arguments are wrong (too many, unknown names, given twice, full without a value, a negative std, lo more than hi); normal or uniform fills a field that isn’t float; or the values it gives don’t fit the field’s domain.

dim D = 64

model Mlp:
    param w1: [D D] = gaussian(std=0.02)

See Initializers.

E0904

Value isn’t the expected model. A parameter or return type is a model, and the value is another model or a tensor. Models are matched by name.

dim D = 64

model A:
    param w: [D D] = zeros

model B:
    param w: [D D] = zeros

fn apply(m: A, x: [b D]) -> [b D] = x @ m.w

fn f(m: B, x: [b D]) -> [b D] = apply(m, x)

See Using models.

10: Devices

E1001

Invalid devices declaration. The device isn’t cpu, mps or cuda; cuda has no device numbers, or a number twice; or cpu or mps has numbers.

devices tpu

See Declaring devices.

E1002

Unknown devices option. An indented line under a devices declaration is an option, for running on several devices. The options are shard params and split NAME, one per line; any other line, shard without params, and split without one name are errors.

devices cuda[0 1]
    shard parameters

Write shard params or split H. See Options for several devices.

E1003

Devices declared more than once. A file has a second devices line, or a notebook cell declares devices other than the ones an earlier cell declared. A program declares its devices once; in a notebook, restart the kernel to change them.

devices cpu
devices cuda[0]

See Declaring devices.

E1004

devices after other declarations. The devices line comes before every other declaration in a file or cell. Comments and blank lines may come before it.

dim D = 64

devices cpu

See Declaring devices.

E1005

Standalone build without devices. shapescript build --standalone makes a module that doesn’t use the shapescript package, so it can’t ask which devices the program declared. The file needs a devices line, or the build needs --devices.

$ cat model.ss
fn double(x: [n]) -> [n] = x * 2
$ shapescript build --standalone model.ss -o model.py
error[E1005]: a standalone build needs to know its devices

See Building for a device.

E1006

--devices disagrees with the file’s devices. shapescript build --devices, or devices= in shapescript.compile, says different devices than the file’s devices line. Remove one of them, or make them agree.

$ cat model.ss
devices cuda[0]

fn double(x: [n]) -> [n] = x * 2
$ shapescript build --devices "cuda[1]" model.ss -o model.py
error[E1006]: the file declares `devices cuda[0]`, but `--devices` says `cuda[1]`

See Building for a device.

11: mut and gradients

E1101

mut where it isn’t allowed. A parameter that isn’t a model or optimizer state, or a parameter of an extern fn, is declared mut; mut marks an argument whose parameter isn’t mut, an argument of a built-in or of Mlp.init, or the gradient of update; or what follows mut isn’t a name.

dim D = 64

fn double(mut x: [D]) -> [D] = x * 2

See mut parameters.

E1102

Change to something not passed as mut. A function assigns a field of, or passes as mut, a model or optimizer state it may not change: a parameter without mut, a value that may be one, or a Python variable.

dim D = 64

model Mlp:
    param w: [D D] = zeros

fn reset(m: Mlp):
    m.w = m.w * 0

See What counts as a change.

E1103

Argument for a mut parameter without mut. The function changes this parameter, so every call writes mut before its argument, as in reset(mut m). update changes its model and optimizer state, so it’s written update(mut m, mut opt, g).

dim D = 64

model Mlp:
    param w: [D D] = zeros

fn reset(mut m: Mlp):
    m.w = m.w * 0

fn f(mut m: Mlp):
    reset(m)

See mut parameters.

E1104

New value for a mut parameter. A mut parameter is assigned a new model or new optimizer state, or used as a loop variable. The caller would still have what it passed, so change the parameter’s fields, or update the state with update, instead.

dim D = 64

model Mlp:
    param w: [D D] = zeros

fn reset(mut m: Mlp):
    m = Mlp.init(seed=0)

See What counts as a change.

E1105

Function that can’t be differentiated. The first argument of value_and_grad isn’t the name of a fn, or the fn’s first parameter isn’t a model, or it has a mut parameter.

dim D = 64

model Mlp:
    param w: [D D] = zeros

fn loss(x: [b D], m: Mlp) -> [] = sum(x @ m.w)

fn f(m: Mlp, x: [b D]) -> []:
    l, g = value_and_grad(loss, x, m)
    return l

See value_and_grad.

E1106

Gradient of something that isn’t a []. value_and_grad differentiates a function that returns a single float, such as a loss. Reduce a tensor to one number first, with mean or sum.

dim D = 64

model Mlp:
    param w: [D D] = zeros

fn out(m: Mlp, x: [b D]) -> [b D] = x @ m.w

fn f(m: Mlp, x: [b D]) -> [b D]:
    l, g = value_and_grad(out, m, x)
    return l

See value_and_grad.

E1107

Gradient of a function that doesn’t read the model. The differentiated function reads no param of its model, so the gradient would always be zero.

dim D = 64

model Mlp:
    param w: [D D] = zeros

fn loss(m: Mlp, x: [b D]) -> [] = sum(x)

fn f(m: Mlp, x: [b D]) -> []:
    l, g = value_and_grad(loss, m, x)
    return l

See value_and_grad.

E1108

Operation that gradients don’t have. Gradients can be added to and subtracted from gradients of the same model, negated, and multiplied or divided by a [] or a number; norm(g) gives their norm. Any other operation on one is an error.

dim D = 64

model Mlp:
    param w: [D D] = zeros

fn f(g: Grad[Mlp]) -> Grad[Mlp] = g * g

See Gradient types.

W1101

Param that a differentiated function never reads. Its gradient is always zero. Remove the param, or make it a buffer if it isn’t trained.

dim D = 64

model Mlp:
    param w: [D D] = zeros
    param unused: [D] = zeros

fn loss(m: Mlp, x: [b D]) -> [] = sum(x @ m.w)

fn f(m: Mlp, x: [b D]) -> []:
    l, g = value_and_grad(loss, m, x)
    return l

See value_and_grad.

12: Modes and random operations

E1201

Call without its Mode. A call leaves out an argument whose parameter is a Mode, as a call to dropout without one does. There is no global training flag, so whether a call trains, and the seed its random numbers come from, are arguments. Pass Train(seed=…) to train or Eval to evaluate.

fn f(x: [b D]) -> [b D] = x |> dropout(0.1)

See Modes.

E1202

Value isn’t a Mode. A Mode parameter or return type got something other than Train(seed=…), Eval or a value of type Mode, such as a bool that says whether to train.

fn f(x: [b D], training: bool) -> [b D] = dropout(x, 0.1, training)

See Modes.

13: Data records and datasets

E1301

Data record has no such field. x.name reads a field that the batch’s or example’s data record doesn’t declare. The help suggests a field with a similar name.

dim D = 64

data Digit:
    img: [D]

fn f(batch: [b] Digit) -> [b D] = batch.imgs

See Batches.

E1302

Invalid data field. A data record’s field has ?k in its type, has no shape, is a host scalar, a model, a data or logs record, a dataset, a gradient, optimizer state, a Schedule, a Mode or a LogSink, or is named with a Python keyword or a name that starts with _. Every example of a record must have the same shapes, so that examples can be stacked into batches; for values of different lengths, pad them to a dim and add a Mask field.

data Sentence:
    tokens: [?n]

See Data records.

E1303

Value isn’t the expected data. A parameter or return type is a batch of a data record or a dataset, and the value is a batch or dataset of another record, a batch with another number of batch axes, a dataset where a batch is expected (or the other way round), or a tensor. Records are matched by declaration.

dim D = 64

data Digit:
    img: [D]

data Pair:
    a: [D]

fn mean_img(batch: [b] Digit) -> [D] = mean(batch.img, over=b)

fn f(p: [b] Pair) -> [D] = mean_img(p)

See Batches and Datasets.

E1304

Invalid dataset stage. shuffle or batches is given something that isn’t a dataset, as in x |> batches(64) for a tensor x; shuffle comes after batches, which would shuffle whole batches; batches is used twice; or the batch size is known to be less than 1.

dim D = 64

data Digit:
    img: [D]

fn f(train: Dataset[Digit]) -> int:
    n = 0
    for batch in train |> batches(64) |> shuffle(seed=0):
        n = n + 1
    return n

See Datasets.

14: Logs and sinks

E1401

Logs record has no such metric. r.name reads, or Name(name=…) gives, a metric that the logs record doesn’t declare. The help suggests a metric with a similar name.

logs StepLogs:
    loss: [] mean

fn step(x: [n]) -> StepLogs = StepLogs(los=sum(x))

See Records.

E1402

Invalid logs metric. A metric of a logs declaration isn’t a single number (a [] tensor or an int, float or bool); doesn’t end with how its values combine, or ends with a word other than mean, sum, max, min or last; or its name is a Python keyword or starts with _.

dim D = 64

logs StepLogs:
    loss: [D] mean

See Logs declarations.

E1403

Value isn’t the expected logs record. A parameter or return type is a logs record and the value is another record or something else, or sink.write is given something other than a logs record. Records are matched by name.

logs StepLogs:
    loss: [] mean

fn report(sink: LogSink, x: [n]):
    sink.write(sum(x))

See Records and Sinks.

E1404

Value isn’t a LogSink. A LogSink parameter or return type got a tensor, a model or a logs record. A sink is a Python object with a write(dict) method, passed in from Python.

logs StepLogs:
    loss: [] mean

fn report(sink: LogSink, l: []) = sink.write(StepLogs(loss=l))

fn f(x: [n]) = report(x, sum(x))

See Sinks.

15: Optimizers, gradient transforms and schedules

E1501

Optimizer that isn’t given a model. sgd and adamw make optimizer state for a model, given first, by position. The state’s type comes from the model’s, so the model can’t be a tensor, a Python value or missing.

dim D = 64

fn f(x: [b D]):
    opt = adamw(x, lr=3e-4)

See Optimizer state.

E1502

Optimizer state or gradient of another model. update was given optimizer state or a gradient of another model type than the model it updates, or a parameter or return type that is optimizer state, such as AdamW[Mlp], got state for another model.

dim D = 64

model Mlp:
    param w: [D D] = zeros

model Other:
    param w: [D D] = zeros

fn step(mut m: Mlp, mut opt: AdamW[Other], g: Grad[Mlp]):
    update(mut m, mut opt, g)

See update and Optimizer state.

E1503

Value isn’t the expected optimizer state. Optimizer state was expected, by update or by a parameter or return type such as AdamW[Mlp], and the value is something else: a tensor, a model, or the state of another optimizer, such as SGD[Mlp] for AdamW[Mlp].

dim D = 64

model Mlp:
    param w: [D D] = zeros

fn step(mut m: Mlp, mut opt: SGD[Mlp], g: Grad[Mlp]):
    update(mut m, mut opt, g)

fn f(mut m: Mlp, g: Grad[Mlp]):
    opt = adamw(m, lr=3e-4)
    step(mut m, mut opt, g)

See Optimizer state.

E1504

Learning rate that isn’t a number or a Schedule. An optimizer’s lr=, warmup’s then, or a Schedule parameter or return type got something other than a Python number or a schedule, such as a tensor. A learning rate that changes is a schedule, as in lr=cosine(3e-4, 1000).

dim D = 64

model Mlp:
    param w: [D D] = zeros

fn f(m: Mlp, lr: []):
    opt = adamw(m, lr=lr)

See Schedules.

E1505

update or a gradient transform given the wrong kind of value. update’s first argument isn’t a model or its third isn’t a gradient, or clip_norm wasn’t given a gradient first.

dim D = 64

fn f(x: [b D]):
    y = clip_norm(x, 1.0)

See update and Gradient transforms.

16: Checkpoints

E1601

Value that save can’t write to a checkpoint. A checkpoint holds models, optimizer state and int, float and bool values. A tensor, a gradient, a Mode, a Schedule, a logs record, a dataset or several results can’t be saved. A tensor that training needs belongs in a model, as a buffer if it isn’t trained; a schedule is saved with the optimizer state that uses it.

dim D = 64

model Mlp:
    param w: [D D] = zeros

fn f(m: Mlp, x: [b D]):
    save("ckpt.pt", m, x)

See save.

E1602

load given something that isn’t a type it can read. After the path, load takes the type of each value in the file: a model’s name, optimizer state such as AdamW[Mlp], or int, float or bool. A value, as in load("ckpt.pt", m), another kind of type, such as a tensor type, Grad[Mlp] or Schedule, or an expression that isn’t a type is this error.

dim D = 64

model Mlp:
    param w: [D D] = zeros

fn f(m: Mlp) -> Mlp:
    return load("ckpt.pt", m)

Write the type: load("ckpt.pt", Mlp). See load.

E1603

Checkpoint path that isn’t a string. save and load take the file’s path first, by position: a string, or a Python value, which is checked when the call runs. Anything else there, such as a model or a tensor, or a path given by name, is this error.

dim D = 64

model Mlp:
    param w: [D D] = zeros

fn f(m: Mlp):
    save(m)

Write save("ckpt.pt", m). See save.

17: Export

E1701

Python in an exported function. An export fn, and every function it calls, can’t use a name from import python, or anything reached through one: a call such as torch.load("w.pt"), a value such as torch.pi, or a pipe stage such as |> torch.tanh. In a notebook, a Python variable is the same. torch.export records torch operations only, and a Python call could do anything.

dim D = 768

import python torch

export fn load_qkv() -> [D 3*D]:
    return torch.load("w_qkv.pt") as [D 3*D]

Pass the value in as a parameter, or give a torch function a type with extern fn, as in extern fn erfc(x: Real) -> NonNeg = "torch.special.erfc". See Export functions.

E1702

extern fn outside torch, called by an exported function. An exported function can call an extern fn whose Python function is in torch or one of its submodules, or in math. One from another module could run any Python, which torch.export can’t record.

extern fn mystery(x: [...]) -> [...] = "mylib.mystery"

export fn f(x: [n]) -> [n]:
    return mystery(x)

See Export functions.

E1703

as on a Python value in an exported function. In an exported function every value must come from torch operations. A Python value, such as one an earlier line got from Python (itself E1701), can’t be given a type with as there.

dim D = 768

import python numpy

export fn scale(x: [n D]) -> [n D]:
    w = numpy.ones(768)
    return x * (w as [D])

Pass the value in as a parameter of the function. See Export functions.

E1704

Value or operation with no torch-only lowering in an exported function. Models, modes, logs records, sinks, gradients, optimizer state, schedules, datasets and batches exist only on the host, and so do the built-ins that make or use them, such as print, save, value_and_grad and update. Changing a model’s field and a for loop over anything but a range are the same.

dim D = 64

export fn f(x: [n D]) -> [n D]:
    print(x)
    return x

An exported program is for inference: pass the model’s fields in as tensors. See Export functions.

E1705

Parameter of an export fn that isn’t a tensor with a shape. The parameters of an exported function are the program’s inputs, and torch.export traces it with example tensors whose shapes, and which of whose axes may change, come from the parameters’ types. A parameter without a shape, a model, or an int, float or bool without a default has none.

export fn f(x: [n], k: int) -> [n]:
    return x * k

Give it a tensor type such as [b n D], or a host scalar a default, which the exported program keeps. See Export functions.

E1706

return inside an if on a tensor in an exported function. An if on a [] Mask becomes torch.cond, whose branches each give values, which come out of the if; a branch can’t return from the function.

export fn f(x: [n]) -> [n]:
    if sum(x) > 0:
        return x
    return -x

Assign the value in both branches and return it after the if: if sum(x) > 0: y = x / else: y = -x, then return y. See Control flow in exported programs.

18: Several devices

E1801

Operation across the batch axis with no lowering for several devices. With several devices each one holds part of the batch, so an operation that combines values across the batch axis needs each device’s part. Sums, means, cross_entropy, and @, matmul, einsum and reduce with sum that contract the batch axis have a distributed lowering; any other reduction over it, an operation along it such as softmax or sort, indexing or slicing it, and causal with it don’t, yet.

devices cuda[0 1 2]

dim D = 4

data Ex:
    x: [D]

fn largest(batch: [b] Ex) -> [D]:
    return max(batch.x, over=b)

Compute it per example, or with one of the operations that have a lowering, or declare one device. See Operations across the batch.

E1802

Batch axis that can’t stay split across the devices. With several devices, a value with the batch axis is split along it, each device holding some of the examples, so the batch axis has to stay an axis of its own, with a generic size that each device’s part gives. Grouping it with another axis, as (b d), a value with the batch axis twice, a batch whose size is fixed, the batch’s size used as a number (on each device, it would be that device’s part’s), and a value with the batch axis passed to a parameter without a shape are errors.

devices cuda[0 1]

dim D = 4

data Ex:
    x: [D]

fn flat(batch: [b] Ex) -> [b*D]:
    return rearrange(batch.x, b d -> (b d))

Keep the batch axis apart, as in b (h w) rather than (b h) w, and use mean rather than a sum divided by b. See The batch axis.

E1803

Device work in a function that runs in the controller. With several devices, a function that loops over a dataset, prints, saves, loads, writes to a sink or uses Python runs in the controller, and so does a function that calls one. value_and_grad, update and changing a model’s field are device work, which runs on the workers, so they can’t be in such a function.

devices cuda[0 1 2]

dim D = 4
dim C = 3

model M:
    param w: [D C] = normal(std=0.02)

data Ex:
    x: [D]
    y: [] Index[C]

fn loss(m: M, batch: [b] Ex) -> []:
    return cross_entropy(batch.x @ m.w, batch.y, over=C)

fn fit(mut m: M, mut opt: AdamW[M], data: Dataset[Ex]):
    for batch in data |> batches(64):
        l, g = value_and_grad(loss, m, batch)
        update(mut m, mut opt, g)

Move the device work into a function of its own, as the README’s train_step is, and call it from the loop. See Where functions run.

E1804

Whole value passed where a function takes a split one. A function on the devices is compiled once. If one call gives it a value with the batch axis, its parameter is split across the devices, and every call must give it one; a value without the batch axis is whole on every device.

devices cuda[0 1]

dim D = 4

model M:
    param w: [D D] = normal(std=0.02)

data Ex:
    x: [D]

fn double(x: [n D]) -> [n D] = x * 2

fn both(m: M, batch: [b] Ex) -> [D D]:
    a = double(batch.x)
    return double(m.w)

Call a function of its own for whole values. See The batch axis.

E1805

Standalone build for several devices. With several devices, the program runs in a controller that starts a worker process for each device, which is the shapescript package’s runtime, so a module built with --standalone, which doesn’t use the package, can’t run on them.

$ cat model.ss
devices cuda[0 1 2]

fn double(x: [n]) -> [n] = x * 2
$ shapescript build --standalone model.ss -o model.py
error[E1805]: a standalone build can't run on 3 devices

Build without --standalone, or for one device. See Building for a device.

19: Python calls and options on several devices

E1901

Python call given data split across the devices. With several devices, Python calls run once, in the controller, as they would on one device. A batch is split across the devices, each holding some of the examples, and so is every value with its batch axis; giving one to a Python function would gather the batch through the controller’s CPU at every step.

devices cuda[0 1 2]

import python augment

dim D = 4

data Ex:
    x: [D]

fn augmented(batch: [b] Ex) -> [b D]:
    return augment.random_crop(batch.x) as [b D]

Move the call into the dataset, so it runs on each example as it’s loaded, or declare the function with extern fn, as extern fn random_crop(x: [b D]) -> [b D] = "augment.random_crop", so each device runs it on its own examples. See Python calls on several devices.

E1902

Split axis whose size the devices don’t divide. split H gives every device the same share of each axis of size H, so the number of devices must divide H.

devices cuda[0 1 2 3 4]
    split H

dim H = 12

Use a number of devices that divides it, here 2, 3, 4, 6 or 12, or split another dim. See Options for several devices.

E1903

devices option given twice. Each option is given once. The devices are in a row, so they can split values along one axis only: a second split, of another dim, is an error too, since a value with both axes couldn’t be split along both.

devices cuda[0 1]
    split H
    split D

dim D = 8
dim H = 2

Keep one split. See Options for several devices.

E1904

split of something other than a dim. split NAME splits every axis whose size is the dim NAME, so NAME must be a dim the file declares: an axis name, such as a generic b, isn’t one.

devices cuda[0 1]
    split heads

fn f(x: [heads n]) -> [heads n] = x * 2

Declare the size as a dim, as dim H = 12, and split that. See Options for several devices.

20: Fixed shapes and padding

H2001

The last batch isn’t padded. The loader pads the last batch of a loop over batches(n) to n with masked-out examples, so that every step has the same shapes, only when every operation across the batch axis has a masked version. Here one doesn’t, so the last batch keeps its own size. The program is correct either way: this is a hint, not a warning.

data Ex:
    x: [4]

fn ranked(batch: [b] Ex) -> [b 4]:
    return sort(batch.x, over=b)

fn epoch(train: Dataset[Ex]):
    for batch in train |> batches(8):
        print(mean(ranked(batch)))

hint[H2001]: the last batch isn't padded to the batch size: `sort` over `b` has no masked version

To pad it, compute what needs the operation per example, or outside the loop. See Fixed shapes and padding, which lists the operations with masked versions.

22: Grouped sizes and shape aliases

E2201

Shape alias defined in terms of itself. A shape alias stands for a fixed list of axes, so it can’t contain itself, directly or through other aliases. Each alias in the cycle is reported.

dim Ping = [Pong n]
dim Pong = [Ping m]

Write the axes out in one of them. See Shape aliases.

E2202

Shape alias where its axes can’t be spliced. A shape alias stands for several axes, so it can be a type’s whole shape or among a shape’s axes, but not part of one axis: not inside a group, in arithmetic, or as an Index[V] size.

dim Images = [b 3 16gh 16gw]

fn f(x: [(Images 2)]) -> [b] = sum(x, over=b)

Put the alias among the axes, as in [Images 2], or write the group out. See Shape aliases.

E2203

Grouped size like 16gh outside an axis. A grouped size is one axis of a type or a pattern, the group (gh 16). It isn’t a size that arithmetic can use, a dim’s value, or an expression, and its name can’t be a keyword.

fn f(x: [b 16gh*2]) -> [b] = sum(x, over=b)

For the product as a size, write 16*gh; for a value, 16 * gh. See Grouped sizes.

H2201

16 gh, two axes, where 16gh was probably meant. In a function’s signature, a number greater than 1 directly followed by a name that isn’t a dim, is an axis of its own nowhere else in the signature, and is part of a group, other than a grouped size, somewhere in it. 16 gh is two axes; 16gh, without the space, is one, of size 16 × gh. This is a hint, not a warning: both are valid.

dim D = 768

fn patchify(img: [b 3 16 gh 16 gw]) -> [b (gh gw) D]:
    return img |> rearrange(b c p gh q gw -> b (gh gw) (p q c))

hint[H2201]: `16 gh` is two axes: did you mean `16gh`?

If two axes were meant, nothing needs to change. See Grouped sizes.

23: Modules, overloaded functions and generic functions

E2301

No .ss file for an import. import name looks for name.ss next to the importing file, then in each folder of the Python path.

import attentoin

Check the name, or put the folder that has the file on the Python path. See Importing other files.

E2302

Files that import each other. A file can’t import itself, directly or through other files: each file is checked with what it imports, so a cycle has nowhere to start.

$ cat a.ss
import b
$ cat b.ss
import a
$ shapescript check a.ss
error[E2302]: `b.ss` imports this file back
  = note: the files import each other: a.ss -> b.ss -> a.ss

Move what both need into a third file that both import. See Importing other files.

E2303

An imported file has errors. A file is checked with the files it imports, and an error in one of them is reported at the import that leads to it, with its errors as notes.

$ cat broken.ss
fn f(x: [n]) -> [n] = y
$ cat main.ss
import broken
$ shapescript check main.ss
error[E2303]: `broken.ss` has an error
  = note: broken.ss:1:23: error[E0301]: unknown name `y`

Fix the imported file; shapescript check on it shows the errors in full. See Importing other files.

E2304

A name an imported file doesn’t declare. std.name must be declared at the top level of the file imported as std.

import shapescript.std.fit as std

fn f(x: [n]) -> [n] = std.nothing(x)

Check the spelling, and that the declaration is in that file, not in one it imports: a file’s own imports aren’t part of its names. See Importing other files.

E2305

Standalone build of a file with imports. A --standalone module needs only torch, but it would import the other files through the shapescript package’s import hook.

$ cat main.ss
import shapescript.std.fit as std
$ shapescript build --standalone --devices cpu main.ss
error[E2305]: a standalone build can't import other `.ss` files

Build without --standalone, or put what you import into the file. See Importing other files.

E2306

An import or generic function where it can’t be used. Importing other .ss files, and calling generic functions, work in files, not in notebook cells; a program on several devices can’t import them yet, since each device’s worker runs only the file’s own code; an exported program can’t call another file’s functions; and an imported file can’t declare other devices than the file importing it.

%%shapescript
import layers
error[E2306]: `import` of a `.ss` file can't be used here

Put what you import into the cell or the file. See Importing other files.

E2307

Functions of the same name for the same model. Functions may share a name only when each takes a different model as its first parameter, since a call picks one by its first argument’s model.

dim D = 4

model A:
    param w: [D] = zeros

fn forward(m: A, x: [D]) -> [D] = x * m.w

fn forward(m: A, x: [D]) -> [D] = x + m.w

Rename one of them. A function whose first parameter isn’t a model, and a generic function, can’t share their names. See Overloaded functions.

E2308

No function of that name for the argument’s model. Each function of the name takes a model first, and none takes the first argument’s.

dim D = 4

model A:
    param w: [D] = zeros

model B:
    param v: [D] = zeros

model Z:
    param z: [D] = zeros

fn forward(m: A, x: [D]) -> [D] = x * m.w

fn forward(m: B, x: [D]) -> [D] = x + m.v

fn f(z: Z, x: [D]) -> [D] = forward(z, x)

Declare a forward for Z, in this file or in the one that declares Z. See Overloaded functions.

E2309

Misused type variable. A type variable of a generic function can’t also name an axis of its signature, and it must be used by a parameter, since calls work it out from their arguments.

dim D = 4

fn f(m: M, x: [M D]) -> [D] = x[0]

Use another letter for the axis, or a name, such as n. See Generic functions.

E2310

Arguments disagree about a type variable. Each type variable stands for one type in a call, worked out from the arguments: two arguments that make it different types, or no argument that says what it is, are an error.

dim D = 4

model A:
    param w: [D] = zeros

model B:
    param v: [D] = zeros

fn same(a: M, b: M) -> [D] = zeros_like(a)

fn f(a: A, b: B) -> [D] = same(a, b)

Pass values of the same type, or give the function two type variables, as fn same(a: M, b: N). A value from Python has no ShapeScript type to give. See Generic functions.

E2311

A generic function needs a function the model doesn’t have. A generic function’s instance for a model calls a function on it that isn’t declared for that model, in the generic function’s file or the model’s. The error is at the call to the generic function, and names where it uses the function.

import shapescript.std.fit as std

dim D = 4

model Plain:
    param w: [D] = zeros

data Ex:
    x: [D]

fn train(mut m: Plain, data: Dataset[Ex], sink: LogSink, mode: Mode):
    std.fit(mut m, data, sink, 1, 8, 0.1, mode)

error[E2311]: `fit` with M = Plain, E = Ex needs a function `loss` that takes `Plain` first
  = note: `train_step` uses it at <shapescript>/std/fit.ss:29:27, called by `run` …

Declare it, here fn loss(m: Plain, batch: [b] Ex, mode: Mode) -> [], next to the model. See Generic functions.

E2312

A generic function doesn’t check with these types. Checking the instance of a generic function for the types of a call found an error. The error is at the call, with the error in the generic function as a note.

dim D = 4

model A:
    param w: [D] = zeros

fn loss(m: A) -> [D] = m.w

fn grad_of(m: M) -> Grad[M]:
    l, g = value_and_grad(loss, m)
    return g

fn f(m: A) -> Grad[A] = grad_of(m)

Here loss for A doesn’t give a [], so it can’t be differentiated. Change the function for the type, or call the generic function with another. See Generic functions.

E2313

Generic functions call each other too deeply. Each level of generic functions that call others for new types makes a round of new instances, and there may be 16.

dim D = 4

model A:
    param w: [D] = zeros

fn g1(m: M) -> [D] = g2(m)
fn g2(m: M) -> [D] = g3(m)
fn g3(m: M) -> [D] = g4(m)
fn g4(m: M) -> [D] = g5(m)
fn g5(m: M) -> [D] = g6(m)
fn g6(m: M) -> [D] = g7(m)
fn g7(m: M) -> [D] = g8(m)
fn g8(m: M) -> [D] = g9(m)
fn g9(m: M) -> [D] = g10(m)
fn g10(m: M) -> [D] = g11(m)
fn g11(m: M) -> [D] = g12(m)
fn g12(m: M) -> [D] = g13(m)
fn g13(m: M) -> [D] = g14(m)
fn g14(m: M) -> [D] = g15(m)
fn g15(m: M) -> [D] = g16(m)
fn g16(m: M) -> [D] = g17(m)
fn g17(m: M) -> [D] = g18(m)
fn g18(m: M) -> [D] = m.w

fn f(a: A) -> [D] = g1(a)

Make the chain shorter, or make some of its functions take the concrete type. See Generic functions.

24: Iterator chains: what they iterate and their independence (comprehensions before CTL-3)

E2401

Iteration over something that isn’t an axis or a tensor. An iterator chain starts with a size (an axis name, a dim, a positive number, or a sum or product of them), a range whose ends are sizes, or a tensor with at least one axis, because the result’s axis is the one iterated and its size must be known from the types. An int known only when the program runs, a [] tensor or a Python value isn’t one. (Until CTL-3 this was the code of a comprehension over something other than an axis.)

dim D = 4

fn firsts(x: [n D], k: int) -> [n D] = k |> map(i => x[0])

Iterate the axis, n |> map(i => x[i]), and slice the result to keep part of it; or use a for loop to repeat something a run-time number of times. See What a chain iterates.

E2402

Iterator stage that reads the result it builds. An iterator stage computes every element at once, so no element can read another; a lambda that uses the name the chain is being bound to, when nothing had that name before, would read other elements of the result.

dim D = 4

fn running(x: [n D]) -> [n D]:
    y = n |> map(i => y[i] * 0.5 + x[i])
    return y

Each element computed from the one before is a recurrence: compute it with a for loop, which runs in order. See Independence.

E2403

Iterator stage whose function does more than compute a value. The function of a map or filter runs for every element at once, so it may only compute a value: no mut arguments, Python calls, print, checkpoints, modes (random draws), value_and_grad, optimizers, schedules, datasets, logs or sinks, no slice bounds that depend on a parameter, and no function that changes a model, branches on a tensor or loops over a Python value. A map’s function must give a tensor.

dim D = 4

fn noisy(x: [n D], mode: Mode) -> [n D] = x |> map(row => dropout(row, 0.1, mode), over=n)

Do it before or after the chain: here, dropout(x, 0.1, mode) on the whole tensor. See Independence.

E2404

Iterator stage that uses values split across devices. With several devices, a value with the batch axis, or with split H a value with an axis of size H, is split across them, each holding a part. An iterator stage in a function on the workers has no lowering that works on the parts, so it can’t use such a value.

devices cuda[0 1]

dim D = 4

data Ex:
    x: [D]

fn doubled(batch: [b] Ex) -> [b D] = batch.x |> map(row => row * 2.0, over=b)

Compute it with operations on whole tensors, such as batch.x * 2.0, which the checker splits, or declare one device. See How iterator chains run.

H2401

Sequential loop whose iterations are independent. A for i in 0..n: over a whole axis whose iterations only add terms that don’t depend on each other to accumulators, as in acc = acc + f(x[i]), could compute every term at once. The loop is correct as it is: this is a hint.