The Editor

ShapeScript was designed for the editor as much as for the command line. Its language server, shapescript-lsp, gives every LSP editor errors as you type, hovers, navigation and more, and the VS Code extension adds pictures: shape diagrams, activation curves and model architectures, drawn as SVG. Chapter 1 shows how to set up VS Code: installing the extension, which is coming to the Marketplace, and selecting your project’s Python. If something there doesn’t work, see Appendix D.

This chapter tours what the editor shows. Open examples/attention.ss and try each one as you read.

Errors as you type

Every diagnostic in this book appears in the editor as you type, underlined where the checker points. Hovering the underline shows the message, and its code links to the reference. Fix it and the underline goes.

Hover: shapes as diagrams

Hover any expression, any pipe stage or any name, and you see its type, drawn:

• 1-, 2- and 3-axis tensors are a bar, a rectangle or a box, with each side labeled.
• More axes are a stack of copies of the last two, labeled with the rest (× H × b).
• Grouped axes, such as (H Dh), are drawn as segmented edges.
• Masks show their structure. A causal mask is drawn as a triangle, in VS Code as SVG and elsewhere as text:

mask: [n m] Mask  ·  lower-triangular  ·  m ≡ n
         m
    ◄─────────►
  ▲ ┌─────────┐
  │ │█        │
  n │███      │
  │ │█████    │
  ▼ │███████  │
    └─────────┘

Hovering a dim shows its value and doc comment, then where it’s used: the values whose types mention it, grouped by function. Hovering an axis name shows where it was introduced. Hovering a built-in that moves values, such as rearrange or split, explains it, with an example, and then shows this call’s tensors before and the result after, side by side. Hovering |>, or -> in a pattern, does the same for that step.

Hover: what a function does

Hover a function’s name, or a call such as argmax(z, over=C), and the editor runs it on a small example and draws what happens. The example comes from the types, with the dims made small, so your functions get pictures as the built-ins do. When the result’s type says something about its values, such as Prob over m or Index[C], the values going in and coming out are drawn: argmax puts a 1 where each row is largest, and softmax makes rows that add up to 1. Otherwise, the picture shows what depends on what. One value going in is changed, and the values coming out that change are outlined. For the causal_attention of the Introduction, VS Code draws it as an image; other editors show it as text:

x: [b n D]        [b n D]
  D                 D
  · · · · · ·       · · · · · ·
  · · · · · ·       · · · · · ·
n · · ● · · · ──→ n ● ● ● ● ● ●
  · · · · · ·       ● ● ● ● ● ●
  · · · · · ·       ● ● ● ● ● ●
at b = 0          at b = 0

Changing position 2 of x changes the output at positions 2 and after, but not before: the causal mask, seen from outside. The weights w_qkv and w_o aren’t drawn. The editor runs the function in your Python interpreter, which needs PyTorch, only in a trusted workspace, and only if the function doesn’t print, save, load, write logs, call Python or change a mut parameter.

Shapes at the end of each line

In a pipeline written one stage per line, each line ends with its shape, as an inlay hint, updated as you type. Hover the hint for the diagram.

Activation curves

Hover an activation, such as gelu, relu, silu or sigmoid, and you see its formula, its curve with the derivative dashed, and its range. It works for your own functions too, wherever they’re declared, called or piped into: an elementwise function, or any function of one tensor and some numbers that only does elementwise maths to the tensor:

## Squared ReLU, scaled.
fn relu2(x: [b D], k: float = 1.0) -> [b D] = square(relu(x)) * k

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

fn layer(h: [b 64]) -> [b 64]:
    return h |> relu2(0.5) |> swish

Hovering relu2 in layer plots it with k = 0.5, the number written at the call. VS Code draws the curve as SVG; other editors draw it in braille.

Math in comments

Comments can hold LaTeX: $...$ inline, $$...$$ on lines of its own. Hovering shows it typeset in VS Code, and as Unicode text elsewhere; math in a ## comment also appears in the hover of the name it documents. Typing \ offers the supported commands.

## Root-mean-square normalization:
## $$\RMSNorm(x) = \frac{x}{\sqrt{\epsilon + \frac{1}{D}\sum_{i=1}^{D} x_i^2}} \odot g$$
fn rms_norm(x: [b n D], g: [D], eps = 1e-6) -> [b n D]:
    # $\epsilon$ keeps the denominator away from 0.
    return x / sqrt(mean(x * x, over=D) + eps) * g

Math never causes an error. A command outside the supported subset is shown as written, with a hint:

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

A model’s architecture

Click a model’s name, where it’s declared or used as a type, and VS Code draws its architecture in an Architecture tab beside the code, which follows the cursor to other models (the CodeLens Show architecture above each model line opens it again after you close it). Hovering the name shows a one-line summary, such as Embedding → Transformer block × 6 → LayerNorm → Linear · tied. The architecture is read from the functions that take it first: its forward, then its loss. It finds the layers however the code writes them: linear layers, MLPs, attention, embeddings, norms, convolutions, dropout and losses, each box labeled with the type of its result and the params it holds, with residual connections drawn as arrows. A loop over layers is drawn once, marked × L. Hover Gpt in examples/gpt.ss to see a whole transformer, or the model of Chapter 17. Other editors get the same as a list:

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

In the tab, Every operation draws each operation instead of layers, and clicking a box selects the code it came from; the tab zooms and pans for big models. Other editors show the drawing in the hover.

Completion

The editor completes names as you type, everywhere a name can go, and puts first the ones whose types fit. Inside mix(, the values that are [b n D], which mix takes, come first, the first one already selected; then those that would also fit; then the rest, with what wouldn’t type-check last. Each item says what it is, as y: [b n D] or fn softmax(x: [... n] Logit over n, over: Axis[n]) -> [... n] Prob over n, and shows the same documentation as hovering it. Choosing a function writes the call, with its arguments to fill in: softmax(x, over=n), then Tab to go from one to the next.

Some places know more:

• After over=, the list is exactly the axes of the tensor the call works on: for a [b n D], b, n and D.
• After |>, it’s the functions that can take the value piped in, as stages: softmax(over=…), offering the value’s axes, or gelu alone. Functions of yours that take exactly that shape come before the built-ins.
• In rearrange(x, and the other pattern operations, it’s the input’s axes, all of them first, then one at a time; after ->, the names on the left.
• In a type, it’s the dims, the axes the signature already names, and after the shape, domains and dtypes. After as, the value’s own type comes first.
• After ., it’s what’s there: a batch’s fields with their types, an imported file’s declarations, init after a model’s name, and a Python module’s names with their signatures (your own modules are read, never run). After m., a model’s fields and the functions that take it, as calls (forward(m, …)).
• After import , the .ss files you can import; after import python , the Python modules.
• At the start of a line, the keywords that can go there: fn, model, data and the other declarations at the top level, written out with their blocks, and for, if, else and return in a function.

Matching is fuzzy: lnorm finds layer_norm. A name from another file of your project that isn’t imported yet is offered too, after the others, as metrics.accuracy: choosing it adds import metrics at the top.

While you type the arguments of a call, its signature shows above the cursor with the parameter you’re on highlighted, its type included. For overloaded functions, all of them are listed, the one your first argument picks first.

Navigation

• Go to definition works on axes, dims, functions, models, fields and imports, across files. On a built-in such as gelu, it opens the read-only declaration from the prelude. On torch.load, an import python module or an extern fn’s path, it opens the Python source.
• Find references and rename work on axes, dims, domains and functions. Renaming an axis changes it in the signature, the patterns and the expressions; a rename that would change what a name means is refused.
• Completion after m., for a model m, offers its fields and every function that takes the model first, even from files not imported yet. Choosing m.forward writes forward(m, …), and adds the import if one is needed.

Python editors know your .ss files

The extension writes a Python module beside each .ss file that checks, model.py beside model.ss, with typed declarations of its dims, models and functions and their docs. So in a Python file, Pylance and mypy resolve from model import forward, complete it, and show its ShapeScript signature on hover. (shapescript pymodules *.ss writes them from the command line.) A .py file you wrote yourself is never replaced.

Debugging

ShapeScript is debugged with Python’s debugger, debugpy, which VS Code and Jupyter use. Breakpoints in .ss files stop, stepping goes from ShapeScript into the Python it calls and back, and variables show their ShapeScript types and sizes:

x = [b=2 D=4] float32 tensor([[1., …

In VS Code, ShapeScript: Debug File in the Command Palette, or in the ▷ menu above the editor, debugs the current .ss file’s fn main(), or a Python program that imports .ss files, with the import hook installed. To keep a configuration, with arguments, add one of the extension’s shapescript launch configurations to .vscode/launch.json (Run › Add Configuration…). shapescript run runs a program without the debugger.

The debugger runs in the Python you selected for the project (Chapter 1), with the debugpy that comes with Microsoft’s Python extension. Without that extension, install debugpy into the environment, pip install debugpy, and name its Python in the setting shapescript.debug.python. The worker processes of several devices (Chapter 16) can’t be debugged yet.

To look at a tensor rather than read it, right-click it in the Variables view and choose View Tensor. It appears in the Tensor section of the Run and Debug view, under Variables and Watch, as a heatmap you can zoom from the whole tensor down to its numbers, arranged by its type: attention weights [b H n m] Prob over m as one heatmap per head, an image [b 3 h w] Unit in color. Beside it are checks of what the type promises, such as no NaN, probabilities that sum to 1, and zeros above the diagonal after a causal mask, each with Show to go to a value that breaks it. The program computes what the view draws on the GPU, so the tensor can be as big as the GPU holds.

Exercises

1. Open examples/attention.ss and hover m in softmax(over=m). Where does the editor say m comes from?

2. In curves.ss above, change relu2’s body to square(relu(x)) * k + mean(x). Does relu2 still get a curve? Why?

3. Fix the math so the hint goes away:

   ## The mean, $\fraq{1}{n}\sum_i x_i$.
   fn avg(x: [n]) -> [] = mean(x, over=n)