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.
| Operation | Shape | Element type |
|---|---|---|
abs, neg, square, floor, ceil, round, sign, relu | Elementwise, 1 argument | Keeps the dtype |
leaky_relu(x, slope), elu(x, alpha) | Elementwise, 1 or 2 arguments | Float |
clamp(x, lo, hi) | Elementwise, 1 to 3 arguments | Promoted (Dtypes) |
scale, add, sub, mul, maximum, minimum | Elementwise, 2 arguments | Promoted (Dtypes) |
div, pow | Elementwise, 2 arguments | Float |
where(cond, a, b) | Elementwise, 3 arguments | cond is a Mask (E0501); promoted; the hull of a’s and b’s ranges |
exp, log, log1p, expm1, sqrt, rsqrt, sin, cos, tanh, sigmoid, erf, gelu, silu, softplus, log_sigmoid | Elementwise prelude stubs | Ranges from range inference |
softmax, log_softmax | Keep the shape; over= required | Take logits over the axis; give Prob and LogProb over it |
cumsum, cumprod, sort, argsort | Keep the shape; over= required | argsort gives Index[n]; sort keeps the values’ domain along n; cumsum of Prob over n along n is in [0, 1] |
sum, prod, max, min, mean, std, var, logsumexp, norm, argmax, argmin, any, all | Remove the over= axis, or every axis without over= | argmax, argmin give Index[n]; sum, prod of booleans or integers give i64; masks count as 0 and 1 |
matmul(x, y) | As @ | |
masked_fill(x, mask, value) | x’s shape; the mask broadcasts to it | Keeps logits logits |
cross_entropy(z, y, over=C) | [] | z must fit Logit over C; y has z’s shape without C, and is integers or Index[C]; gives NonNeg |
causal[n m] | [n m] | Mask (Indexing) |
topk(x, k, over=n) | Two results, n replaced by an axis of size k | Values (a distribution’s domain becomes its range) and Index[n] indices. k is a number, dim, axis name, or a run-time integer (unknown size); a k larger than the axis or negative is E0415 |
cat(a, b, ..., over=n) | The n axes’ sizes add up; the other axes line up by position | Kept if all equal, else the hull of the ranges |
stack(a, b, ..., over=k) | A new first axis of size the count, named k if given | As cat |
gather(x, ids, over=n) | The indices’ shape; ids has x’s axes, except that n is another axis or left out | ids are integers, checked against n if Index (E0507) |
scatter(x, ids, values, over=n) | x’s shape | values has the indices’ shape, or is one number, and x’s dtype (E0502) |
rearrange, reduce, repeat, einsum, split | Pattern operations | |
linear, layer_norm, rms_norm, batch_norm, conv2d, max_pool2d, avg_pool2d, embedding, scaled_dot_product_attention | Prelude stubs (The prelude) | |
dropout(x, p, mode) | x’s shape; a prelude stub | Float, Real; takes a Mode |
print(a, b, ...) | No result | Prints any values (Printing) |
sgd(m, lr=…), adamw(m, lr=…) | Optimizer state, SGD[M] or AdamW[M] | Optimizer state |
update(mut m, mut opt, g) | No result | One optimizer step (update) |
clip_norm(g, max_norm) | g’s type, Grad[M] | Gradient transforms |
cosine(peak, steps, end), warmup(steps, then) | A Schedule | Schedules |
save(path, values…) | No result | Writes a checkpoint (Checkpoints) |
load(path, Types…) | A value of each type | Reads one (Checkpoints) |
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.