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.