Keyboard shortcuts

Press ← or → to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

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.