Models

Model declarations

model Name: followed by an indented block of fields declares a model, a record of the tensors that make up a program’s state:

## A two-layer classifier.
model Mlp:
    param w1: [D Hd] = normal(std=0.02)
    param b1: [Hd] = zeros
    param w2: [Hd C] = uniform(-0.1, 0.1)
    ## Subtracted from every input. Saved with the model, never trained.
    buffer mean: [D] = zeros

A param field is trained. A buffer field is saved with the model but not trained. A model and each of its fields can have a doc comment. A model’s fields have different names (E0302). Models, like functions, can be used anywhere in the file.

Fields

A field’s type has a shape made only of numbers and dims, so every model of a type has the same shapes: a generic name is E0301, and ?k is E0902. For a single number, write []. The other rules are E0902:

• A param holds floats, of any float dtype (f32 unless the type names another), and has no domain, since training can change it to any value.
• A buffer can have any dtype, and a range domain, Index[V] or Mask, but not a semantic domain, Binary or Additive.
• A field can’t be a host scalar, a type without a shape, or a model: models don’t nest. Nor can it be optimizer state or a Schedule.
• A field can’t be named init, a Python keyword, a name that starts with _, or an attribute of torch.nn.Module, such as train, eval, to or half.

Initializers

Every field has an initializer after its =, which gives it its first value:

Arguments are given by position or by name, as in a call. A missing initializer, one that isn’t in the table, the wrong arguments, a negative std and a lo more than hi are E0903. The arguments are Python numbers, and may use dims, as in normal(std=1 / D): floats, except that full takes an int for an integer field and a bool for a bool field (E0502). normal and uniform only fill float fields, and the values an initializer gives must fit the field’s domain (E0903): zeros can’t start a Pos buffer, nor uniform(0, 2) a Unit one.

Using models

A model’s name is a type, as in m: Mlp for a parameter or -> Mlp for a return type. Only a model of that name fits it; another model or a tensor is E0904. A Python value fits, and is checked when the function runs (Runtime checks). A model’s name in any other type, as in [D] Mlp, is E0303.

m.w1 reads a field, and has the field’s type. A field the model doesn’t declare is E0901. A model isn’t a tensor: arithmetic on one, or passing it for a tensor parameter, is E0412.

Mlp.init(seed=s) makes a model. seed is an int, given by name (E0406; E0502 or E0805 for something other than an int). init is all a model’s name can be used for as a value: Mlp alone and Mlp(...) are E0303, and Mlp.other(...) is E0301.

Compiled models

A model compiles to a Python class, a subclass of torch.nn.Module. Mlp.init(seed=0), called from ShapeScript or from Python, returns one, with each param registered as an nn.Parameter and each buffer with register_buffer. So its state_dict() holds exactly the declared fields, and it works with existing PyTorch code. Storing any other tensor or module on it raises AttributeError.

Seeds. There is no global random state. Each field is drawn from a CPU torch.Generator of its own, seeded with the first 8 bytes of the SHA-256 of "{seed}/{field}", where field is the field’s name. So the same seed gives the same values, and adding, removing or reordering fields doesn’t change the others. Random values are drawn as f32 (f64 for an f64 field) and then converted to the field’s dtype, so they are the same on any device. Every field is made on the declared device.

A model’s architecture in the editor

In VS Code, a model’s architecture is drawn in a tab of its own, Architecture: Mlp, beside the editor. Clicking a model’s name, where it’s declared, used as a type (m: Mlp) or in Mlp.init, or the name of a function that takes a model first (fn forward(m: Mlp, …), or a call of one), shows that model there; the tab follows the cursor from model to model without taking the focus. It opens by itself the first time; once closed, it stays closed until Show architecture above a model line, the hover’s Open architecture link, Show Model Architecture in the editor’s context menu or the command palette opens it again (shapescript.architecture.followCursor: false stops it following the cursor). In the tab, the drawing is fitted to the width and can be zoomed (the buttons, Ctrl/Cmd + scroll, + - 0) and dragged; Layers and Every operation switch the level; the functions it was read from are links to their lines; and clicking a box selects the code it was drawn from. The hover then sums the architecture up in one line, from its layers, Architecture of Gpt: Embedding → Transformer block × 6 → LayerNorm → Linear · tied · Open architecture, above the model’s code and docs.

Other editors show the drawing on hover instead. Hovering a model’s name, where it’s declared or used as a type (m: Mlp), draws its architecture, read from the functions that take it as their first parameter, in its file and the files it imports (overloads): its forward, then its loss. The hover starts by naming them, with links to their lines: “Architecture of Mlp, from forward (mlp.ss:26) and loss (mlp.ss:29)”. Without a forward, it uses the other functions that take the model first and return a tensor (at most two, not those that take it mut); with none, it lists the fields, each with a small drawing of its shape.

The drawing goes from top to bottom, one box per layer, as papers and explainer videos draw a model: Embedding V → D, a Transformer block × L = 6 holding LayerNorm, Multi-head attention · H heads · causal, a residual connection ⊕, LayerNorm, MLP D → Dff → D · GELU and another ⊕, then LayerNorm, Linear D → V · tied to wte and the loss, Cross-entropy over V. Under each layer’s name is the type of its result, as the checker worked it out, with sizes named as the code writes them; the function’s inputs come first and its declared return type last. The params a layer holds are inside its box, each a small rectangle with the proportions of its diagram, its name and its type (just the names when they don’t fit); other inputs from the side, such as the labels cross_entropy compares with, are beside it, with an arrow in.

Layers are found from the operations, however the code writes them:

• Linear: a matrix product with a 2-D param, x @ m.w, matmul(m.w), linear(x, m.w, m.b) or an einsum of one input with a 2-D param, with a 1-D param added after it as its bias: Linear D → C. An activation right after it is part of it, Linear D → Dff · GELU; if its weights are an embedding’s, it says so, · tied to wte.
• MLP: linear layers with an activation after each but the last, and maybe dropout between them: MLP D → Dff → D · GELU, · dropout 0.1.
• Attention: scaled_dot_product_attention, or a softmax over a product of two computed values (einsum(q, k: …)), with the projections that make the queries, keys and values, the steps between (splitting heads, scaling, masking, the product with the values, merging heads) and the projection after: Multi-head attention · H heads, with · causal if a causal mask is applied. H is the axis before the last two of its scores or result.
• Embedding (embedding(ids, m.wte)), position embedding (a param of two or more axes added to every example), LayerNorm, RMSNorm, BatchNorm, GroupNorm (the built-ins, or a function named so that does arithmetic, as rms_norm written out does), convolutions (Conv k×k · C → D, or Patch embedding P×P · 3 → D when its stride is its kernel), dropout (a small tag, Dropout 0.1) and losses (Cross-entropy over V, MSE loss).
• A function that’s called is a frame named after it, unless it’s one layer, which is drawn as that layer; one that holds attention, a feed-forward layer and a residual connection is a Transformer block, with the function’s name on the right.
• Reshapes that only move axes (rearrange, split, transpose, casts) are left out, as diagrams usually leave them out, except a rearrange that cuts an image into patches, Patchify 16×16. Other operations stay boxes of their own, with readable names: Mean over n, Softmax over m, × 0.5.

A value that isn’t the line before’s, such as a residual connection’s x + a, is an arrow down the left, into a + circle. Calls that pass the model or its params to another function are followed, three calls deep, with a link to each in the header; a loop’s body is drawn once, as a frame marked × L (× L = 6 when L is a dim), and so are calls of one function in a row on each other’s results (× 3). Functions given only tensors, and forward called from loss, are one box each. Past 32 boxes, the innermost frames become one box each, and then the last boxes are counted instead of drawn. Layers are colored by kind, in the theme’s chart colors: blue for linear layers, MLPs and embeddings, green for activations, purple for norms, orange for attention and softmax, red for the loss.

The operations level draws every operation instead, as the code does it: matmul, add, layer_norm over D, split b n (3 H Dh) → b H n Dh, each with the params it reads to its right, and each function that’s called in its own frame. In VS Code, the setting shapescript.architecture.detail chooses the level (layers or operations), and the hover’s link, Show every operation or Show layers, switches it and shows the hover again. Other editors choose it with the initialization option architectureDetail: "operations", or later with the notification shapescript/setArchitectureDetail with {detail}; a client that runs command: links in hovers can ask for the link with architectureToggle: true, whose command, shapescript.architecture.toggleDetail, is the client’s to run.

A client can show the drawing in a view of its own, as VS Code’s tab does, with the request shapescript/architecture, given {textDocument: {uri}, position} and optionally detail ("layers" or "operations"). It answers null where there’s no model (a model’s name, or a function that takes a model first), and otherwise {model, title, summary, detail, svg, width, height, text, declaration, links, nodes}: summary is the one-line summary; svg is drawn in currentColor and the accents above, with each box in a <g data-node="id">; text is the list below; declaration is the model’s {uri, range}; links are the functions read (kind: "source") and followed ("followed"), each {label, place, kind, uri, range}; and nodes are {id, uri, range}, where each box’s code is. probe: true answers only {model}. A client that says so with architecturePanel: true gets the summary and a link to the command shapescript.architecture.show in the hover, rather than the drawing.

Editors that show images in hovers show the drawing as an image; a hover is short, so drag its bottom edge down to see all of a large model. Other editors get a list, one line per layer, with ↓ where it reads the line above, ⊕ for a sum with a value from further up, + x naming that value, the params a layer holds after its type, those beside it after ←, and frames drawn with ┌ │ └:

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

loss  mlp.ss:29
  • batch.img                            [b D]
  ↓ forward (above)                      [b C] Logit over C
  ↓ Cross-entropy over C                 []
        ← batch.label [b] Index[C]
  ↓ output                               [] NonNeg

and, at the operations level:

forward  mlp.ss:26
  • x                     [b D]
  ↓ matmul                [b Hd]              ← m.w1 [D Hd]
  ↓ gelu                  [b Hd]
  ↓ dropout 0.1           [b Hd]
  ↓ matmul                [b C]               ← m.w2 [Hd C]
  ↓ output                [b C] Logit over C

Only the first branch of an if is drawn. examples/gpt.ss is a GPT whose hover shows its embedding, its transformer block, with the attention and the MLP inside it, once, × L = 6, and its output layer, tied to the embedding.