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.