Broadcasting by Name, @ and Pipes

The spotter has most of its parts now: patchify cuts an image into patches, norm normalizes tokens, mix_patches mixes the patches, and head turns tokens into class scores. What’s missing is the glue: a matrix product that turns each patch into a token, a second MLP that mixes the channels of each token, and a forward function that runs the whole thing in order.

That glue is operators: how +, * and friends line up tensors of different shapes, how @ multiplies matrices, and how |> strings operations together. This chapter covers all three, and ends with a spotter that goes from pixels to scores.

Broadcasting by name

Look again at norm, from Chapter 4:

## Layer normalization over the channels, written out.
fn norm(x: [b n D], eps = 1e-5) -> [b n D]:
    mu = mean(x, over=D)
    var = mean(square(x - mu), over=D)
    return (x - mu) / sqrt(var + eps)

mu is [b n], one mean per token, and x is [b n D]. In PyTorch, x - mu would fail, or worse, quietly line up the wrong axes: you’d write mean(dim=-1, keepdim=True), or mu[..., None], and getting the Nones wrong is a silent bug when sizes happen to match. Here x - mu just works, because ShapeScript lines axes up by name whenever it can. mu’s axes are b and n, x has axes called b and n, so that’s where they go, and D is broadcast.

A few more that need no unsqueeze:

dim D = 64

fn add_mask(scores: [b H n m], mask: [n m]) -> [b H n m] = scores + mask

fn add_bias(x: [b n D], bias: [b D]) -> [b n D] = x + bias

fn scale_rows(x: [b n D], w: [b n]) -> [b n D] = x * w

## Each image minus its own mean brightness.
fn center(img: [b c h w]) -> [b c h w] = img - reduce(img, b c h w -> b, mean)

The compiled Python inserts the right Nones.

The rule has two cases:

1. By name. If every axis of one operand has a name, and all those names are axes of the other operand, they line up by name. The result has the other operand’s axes, in its order. (If both qualify, as with [n m] + [m n], the left operand’s order wins.)
2. From the right. Otherwise the axes line up from the right, as in NumPy. Each pair must have the same size or one of them size 1, and two named axes in a pair must have the same name.

The second case is what catches bugs. Say you want to give each patch a learned position vector, so the spotter knows where each patch was, and you write the parameter with a fresh axis name:

dim N = 16
dim D = 32

fn add_positions(x: [b N D], pos: [k D]) -> [b N D] = x + pos

$ shapescript check positions.ss
error[E0404]: can't line up the axes of `[b N D]` and `[k D]`
 --> positions.ss:4:55
  |
4 | fn add_positions(x: [b N D], pos: [k D]) -> [b N D] = x + pos
  |                                                       -^^^---
  |                                                       |   |
  |                                                       |   this is `[k D]`
  |                                                       this is `[b N D]`
  |
  = note: neither operand's names are all in the other, so the axes line up from the right, and `N` meets `k`
  = help: give them the same names with `rearrange` or `as`

For more information about E0404, run `shapescript explain E0404`.

k isn’t an axis of x, so the axes line up from the right, and N meets k. Even if a caller passes 16 rows, they’re different axes as far as the function knows. The fix is to say what you mean, pos: [N D]: one vector per patch. If two axes really are the same but have different names, say so: Chapter 7 shows as, and Chapter 5 rearrange.

Unnamed axes, such as the 1 in [b 1 S S] or the 3 of an RGB image, line up from the right only:

## Adds a per-channel offset to an RGB image.
fn shift(img: [b 3 h w], offset: [3 1 1]) -> [b 3 h w] = img + offset

Matrix multiplication

The first step of the spotter turns each patch, 16 pixels, into a token of D numbers. That’s a matrix product. x @ y and matmul(x, y) work like torch.matmul: the last axis of x meets the second-to-last of y (or its only axis), and the axes before those broadcast by the rule above. Add this to spotter.ss:

## A token per patch.
fn embed(img: Images, w_e: [P*P D]) -> [b N D] = patchify(img) @ w_e

[b N P*P] @ [P*P D] is [b N D]: the pixels of each patch meet the rows of w_e, and b and N come along. Batched products work the same way, with the leading axes broadcast:

fn per_head(q: [b H n Dh], k: [b H Dh m]) -> [b H n m] = q @ k

The contracted axes must have equal sizes, but not equal names. That’s deliberate: in x @ w, the meeting axes usually have different roles. In embed it’s “the pixels of a patch” meeting “the inputs of the layer”; they’re both P*P here, but they could as well be Din and P*P. A size mismatch is still an error. Write the weights the other way round, as nn.Linear stores them:

dim P = 4
dim D = 32

fn embed(patches: [b n P*P], w_e: [D P*P]) -> [b n D] = patches @ w_e

$ shapescript check bad_embed.ss
error[E0401]: can't multiply `[b n P*P]` by `[D P*P]`
 --> bad_embed.ss:4:57
  |
4 | fn embed(patches: [b n P*P], w_e: [D P*P]) -> [b n D] = patches @ w_e
  |                                                         -------^^^---
  |                                                         |         |
  |                                                         |         this is `[D P*P]`
  |                                                         this is `[b n P*P]`
  |
  = note: the axes `@` contracts differ: `P*P` = 16, not `D` = 32

For more information about E0401, run `shapescript explain E0401`.

Pipes

The second MLP of a mixer block works on each token on its own, mixing its D channels: normalize, widen to Dc, gelu, narrow back to D, and add the input. Written as nested calls, that’s add(x, matmul(gelu(matmul(norm(x), w1)), w2)), which you read inside out. A pipe reads in the order it runs. Add this:

## Hidden width of the MLP that mixes channels.
dim Dc = 64

## Mixes information between the channels of each token.
fn mix_channels(x: [b N D], w1: [D Dc], w2: [Dc D]) -> [b N D]:
    return x |> norm |> matmul(w1) |> gelu |> matmul(w2) |> add(x, _)

a |> f(args) means f(a, args), and a |> f means f(a). A few rules:

• _ puts the value somewhere other than first. |> add(x, _) is add(x, value), the residual connection. A stage may have one _.
• A line that starts with |> continues the line above, so a pipeline can be written one stage per line, as below.
• |> binds more loosely than everything else. x @ w |> gelu is (x @ w) |> gelu. To apply as or to to a whole pipeline, use parentheses: (x |> f) as T.

Every built-in takes the tensor first, so any of them can be a stage, and so can your own functions: norm above, and, in the next listing, embed, mix_patches, mix_channels and head. That’s the whole spotter, one stage per line:

## The class scores of each image.
fn forward(img: Images, w_e: [P*P D], tw1: [N Dt], tw2: [Dt N],
           cw1: [D Dc], cw2: [Dc D], w_h: [D C]) -> [b C]:
    return img
        |> embed(w_e)
        |> mix_patches(tw1, tw2)
        |> mix_channels(cw1, cw2)
        |> norm
        |> head(w_h)

Read it top to bottom: patches become tokens, the tokens mix with each other, then within themselves, and are normalized, averaged and scored. (A real mixer has several blocks; Chapter 9 adds a loop.) The parameter list continues on a second line, since it’s inside parentheses.

When a stage fails, the message says which one, and what its input was. Here mix_channels uses the first weight twice:

dim N = 16
dim D = 32
dim Dc = 64

fn mix_channels(x: [b N D], w1: [D Dc], w2: [Dc D]) -> [b N D]:
    return x
        |> layer_norm(over=D)
        |> matmul(w1)
        |> gelu
        |> matmul(w1)
        |> add(x, _)

$ shapescript check bad_pipe.ss
error[E0401]: can't multiply `[b N Dc] [-0.1701, inf)` by `[D Dc]`
  --> bad_pipe.ss:10:12
   |
 6 |       return x
   |  ____________-
 7 | |         |> layer_norm(over=D)
 8 | |         |> matmul(w1)
 9 | |         |> gelu
   | |_______________- this is `[b N Dc] [-0.1701, inf)`
10 |           |> matmul(w1)
   |              ^^^^^^^--^
   |                     |
   |                     this is `[D Dc]`
   |
   = note: the axes `@` contracts differ: `Dc` = 64, not `D` = 32
   = note: in stage 4 of this pipeline, whose input is `[b N Dc] [-0.1701, inf)`

For more information about E0401, run `shapescript explain E0401`.

The second matrix should have been w2. (The [-0.1701, inf) after the shape is the range of values gelu can give, which the checker tracks; Chapter 8 explains why.)

In the editor, each line of a pipeline ends with its shape, so you can see this before you run anything (Chapter 13).

The spotter so far

spotter.ss now goes from an image to four class scores. It can’t run anything useful yet, since it has no trained weights; the next chapter calls it from Python for the first time.

devices cpu

## Side of each image, in pixels.
dim S = 16
## Side of each square patch.
dim P = 4
## Number of patches in an image.
dim N = (S / P) * (S / P)
## Width of each token.
dim D = 32
## Hidden width of the MLP that mixes patches.
dim Dt = 64
## Hidden width of the MLP that mixes channels.
dim Dc = 64
## Number of classes: square, disk, plus and cross.
dim C = 4

## A batch of grayscale images.
dim Images = [b 1 S S]

## The images that have any ink: a pixel brighter than 0.5.
fn nonblank(img: Images) -> [?k 1 S S]:
    return img[reduce(img, b c h w -> b, max) > 0.5]

## Layer normalization over the channels, written out.
fn norm(x: [b n D], eps = 1e-5) -> [b n D]:
    mu = mean(x, over=D)
    var = mean(square(x - mu), over=D)
    return (x - mu) / sqrt(var + eps)

## The average token of each image.
fn pool(x: [b n D]) -> [b D] = mean(x, over=n)

## The class scores of each image.
fn head(x: [b n D], w: [D C]) -> [b C] = pool(x) @ w

## Cuts each image into P×P patches, one row of pixels per patch.
fn patchify(img: Images) -> [b N P*P]:
    return img |> rearrange(b c (gh p1) (gw p2) -> b (gh gw) (p1 p2 c), p1=P, p2=P)

## Mixes information between the patches.
fn mix_patches(x: [b N D], w1: [N Dt], w2: [Dt N]) -> [b N D]:
    h = einsum(norm(x), w1: b N D, N Dt -> b Dt D) |> gelu
    return x + einsum(h, w2: b Dt D, Dt N -> b N D)

## A token per patch.
fn embed(img: Images, w_e: [P*P D]) -> [b N D] = patchify(img) @ w_e

## Mixes information between the channels of each token.
fn mix_channels(x: [b N D], w1: [D Dc], w2: [Dc D]) -> [b N D]:
    return x |> norm |> matmul(w1) |> gelu |> matmul(w2) |> add(x, _)

## The class scores of each image.
fn forward(img: Images, w_e: [P*P D], tw1: [N Dt], tw2: [Dt N],
           cw1: [D Dc], cw2: [Dc D], w_h: [D C]) -> [b C]:
    return img
        |> embed(w_e)
        |> mix_patches(tw1, tw2)
        |> mix_channels(cw1, cw2)
        |> norm
        |> head(w_h)

Exercises

1. Give embed a bias, b_e: [D], added after the matrix product. Does the + need an unsqueeze? Then add a learned position vector per patch, pos: [N D], so the spotter knows where each patch came from.
2. Real layer norms have a learned gain and shift. Give norm two more parameters, g: [D] and beta: [D], and apply them to the result. Which broadcasting rule lines them up?
3. Rewrite head as a pipeline that starts with x.
4. mix_patches uses einsum so that the weights meet the patch axis N, not the last axis. With the weights stored the other way round, w1: [Dt N] and w2: [N Dt], it’s w1 @ norm(x): a [Dt N] times a [b N D]. Write this version of mix_patches as a pipeline that starts with x. Hint: _.