Pattern Operations

The spotter can normalize tokens, pool them and score them, but it doesn’t have any tokens yet. It needs two things that are both about moving axes around: cutting each image into patches, a token per patch, and mixing information between the patches. Both are reshaping, and reshaping is where shape bugs breed. PyTorch’s view, permute, transpose and reshape take positions and sizes, and a wrong one often still runs. einops fixed much of this with patterns such as "b n (h d) -> b h n d", but they’re strings, checked only when they run, and the names in one pattern mean nothing to the next line.

ShapeScript builds einops’ patterns into the language. Five operations take one: rearrange, reduce, repeat, einsum and split. This chapter uses them to give the spotter its patches and its first mixing layer.

Cutting images into patches

An image is [b 1 S S]: a batch, one channel, and 16 × 16 pixels. We want [b N P*P]: for each image, N = 16 patches, each a row of the P × P = 16 pixels it covers. Add this to spotter.ss:

## 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)

A pattern is inputs -> output, written without quotes. The left side describes the operand’s axes by position, one entry per axis, and each entry binds to one axis of img:

• A name already in scope must have that axis’s size. b is img’s batch axis, so it matches.
• A new name is introduced with the size of the axis it binds to, and stays in scope for the rest of the function. c binds the channel axis, of size 1.
• A number must equal the axis’s size.
• A group (gh p1) binds one axis whose size is the product of its parts. At most one part may be unknown, and the checker solves for it.

Here the groups bind the two axes of size S. A new name can get its size from a keyword argument, as in einops, and the size may be a number or a dim: p1=P says each patch is 4 pixels high, so the checker solves gh = 16 / 4 = 4, and the same for gw. Those are the 4 × 4 patches down and across.

The right side says what the result looks like. Its names must come from the left (except in repeat), and the result’s axes are named after them: the result is [b (gh gw) (p1 p2 c)], which is [b 16 16]. That fits the return type [b N P*P], because the checker compares sizes after working them out: gh × gw = 16 = N, and p1 × p2 × c = 16 = P*P.

Sizes that don’t divide are an error. Try 5-pixel patches:

dim S = 16

fn patchify(img: [b 1 S S]) -> [b ?n ?k]:
    return img |> rearrange(b c (gh p1) (gw p2) -> b (gh gw) (p1 p2 c), p1=5, p2=5)

$ shapescript check bad_patches.ss
error[E0402]: can't split this axis into `(gh p1)`
 --> bad_patches.ss:4:33
  |
4 |     return img |> rearrange(b c (gh p1) (gw p2) -> b (gh gw) (p1 p2 c), p1=5, p2=5)
  |                                 ^^^^^^^ this axis is `S` = 16
  |
  = note: 16 isn't divisible by 5
  = note: in stage 1 of this pipeline, whose input is `[b 1 S S]`

error[E0402]: can't split this axis into `(gw p2)`
 --> bad_patches.ss:4:41
  |
4 |     return img |> rearrange(b c (gh p1) (gw p2) -> b (gh gw) (p1 p2 c), p1=5, p2=5)
  |                                         ^^^^^^^ this axis is `S` = 16
  |
  = note: 16 isn't divisible by 5
  = note: in stage 1 of this pipeline, whose input is `[b 1 S S]`

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

One error for each side of the image. In PyTorch, img.unfold(2, 5, 5) would quietly drop the last pixel of each row and column.

Names are positions, not lookups

The checker catches a classic einops mistake: a pattern that names an axis after a different axis of the operand. Patterns bind by position, so this almost always means the axes were written in the wrong order. The tokens the spotter will mix are [b N D]; here is an attempt to turn them into [b D N] that writes the left side in the order wanted on the right:

dim N = 16
dim D = 32

fn swap(x: [b N D]) -> [b D N]:
    return x |> rearrange(b D N -> b N D)

$ shapescript check transposed.ss
error[E0408]: the pattern calls axis 2 `D`, but `D` is axis 3 of `[b N D]`
 --> transposed.ss:5:29
  |
5 |     return x |> rearrange(b D N -> b N D)
  |            -                ^ axis 2 is `N`
  |            |
  |            this is `[b N D]`
  |
  = note: patterns name axes by position
  = help: to reorder the axes, use `rearrange` first
  = note: in stage 1 of this pipeline, whose input is `[b N D]`

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

The left side describes what x is, b N D, and the right side what you want, b D N.

rearrange

rearrange moves, merges and splits axes. Every name on the left must be on the right, so no data is lost. patchify merges, and the fix above moves:

dim N = 16
dim D = 32

fn swap(x: [b N D]) -> [b D N] = x |> rearrange(b N D -> b D N)

fn flatten(img: [b c h w]) -> [b (c h w)] = rearrange(img, b c h w -> b (c h w))

fn merge_heads(x: [b H n Dh]) -> [b n (H Dh)] = rearrange(x, b H n Dh -> b n (H Dh))

A 1 on the right adds an axis of size 1, as unsqueeze would. That’s how to give a single image a batch axis:

dim S = 16

fn as_batch(img: [1 S S]) -> [1 1 S S] = rearrange(img, c h w -> 1 c h w)

reduce

reduce removes the axes left out of the right side, combining them with sum, mean, max, min or prod. Chapter 3’s nonblank used one: reduce(img, b c h w -> b, max) is the brightest pixel of each image. Groups work here too. The mean brightness of each patch, which tells you where the ink is, is patchify with a reduction instead of a merge:

dim S = 16
dim P = 4
dim N = (S / P) * (S / P)

## The mean brightness of each patch.
fn patch_brightness(img: [b 1 S S]) -> [b N]:
    return img |> reduce(b c (gh p1) (gw p2) -> b (gh gw), mean, p1=P, p2=P)

## Halves each side of an image, keeping the brightest pixel of each 2×2 block.
fn shrink(img: [b 1 S S]) -> [b 1 S/2 S/2]:
    return img |> reduce(b c (h p1) (w p2) -> b c h w, max, p1=2, p2=2)

For a single named axis, the reductions are also functions that take the axis with over=: the spotter’s pool, mean(x, over=n), is reduce(x, b n D -> b D, mean). sum(x) without over= sums everything to a [].

repeat

repeat adds axes, which may be new names given a size as a keyword. Doubling each pixel is the reverse of shrink:

dim S = 16

fn grow(img: [b 1 S S]) -> [b 1 2*S 2*S]:
    return img |> repeat(b c h w -> b c (h p1) (w p2), p1=2, p2=2)

fn copies(img: [b 1 S S]) -> [b 4 S S] = repeat(img, b c h w -> b (c r) h w, r=4)

Mixing the patches with einsum

Now the spotter’s first mixing layer. After patchify and the step that turns each patch into a token (the next chapter writes it), an image is [b N D]: a token of width D for each of the N patches. A token-mixing MLP runs along the patch axis: for each channel separately, it combines the values of all 16 patches, so one patch learns what’s in the others. That’s an MLP whose weights meet the axis N, with a hidden width of its own. Add a dim for it (why it’s 64 comes up below):

## Hidden width of the MLP that mixes patches.
dim Dt = 64

With rearrange, you’d move N to the end, multiply by w1: [N Dt], apply gelu, multiply by w2: [Dt N], and move it back. einsum says which axes meet directly, so there is nothing to move. It takes several operands, written before a :, and one pattern side per operand. Names shared by several operands must have one size, and names left out of the output are summed over:

## 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)

The first einsum sums over N, the axis that the token and w1 share, and gives [b Dt D]: for each channel, Dt hidden values, each a combination of all the patches. The second sums over Dt and gives back [b N D], which is added to x, so the block is residual. norm is the function from Chapter 4.

einsum is the general product: matrix products, batched dot products and outer products are all patterns. Its operands can be named anything, but the pattern’s names say which axes meet. If you write the output in the wrong order, the return type catches it. Say the first half of the MLP is a function of its own:

dim N = 16
dim D = 32
dim Dt = 64

## The hidden values of the token-mixing MLP.
fn hidden(x: [b N D], w1: [N Dt]) -> [b Dt D]:
    return einsum(x, w1: b N D, N Dt -> b D Dt) |> gelu

$ shapescript check bad_mix.ss
error[E0401]: mismatched shape: expected `[b Dt D]`, found `[b D Dt] [-0.1701, inf)`
 --> bad_mix.ss:7:12
  |
7 |     return einsum(x, w1: b N D, N Dt -> b D Dt) |> gelu
  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the return type is `[b Dt D]`
  |
  = note: axis 2: `D` = 32, not `Dt` = 64
  = note: axis 3: `Dt` = 64, not `D` = 32

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

This is why the spotter’s Dt is 64 and not 32, like D. Dims are compared by their values (Chapter 3), so with Dt = 32 the slip would check: [b D Dt] would be [b 32 32] either way. Give different roles different sizes, and the checker can tell them apart. In mix_patches itself, the same slip in the second einsum, -> b D N, would do no harm: x + ... adds a [b N D] to a [b D N], and the next chapter shows that the axes then line up by name, so the sum is right anyway.

split

split cuts one axis into several results. The pattern’s left side has a number that’s missing from the right; that’s how many results there are, and you bind them to that many names. The spotter has no use for it yet, but it could compare the top half of each image with the bottom half:

dim S = 16

## The top half of each image, minus the bottom half.
fn halves(img: [b 1 S S]) -> [b 1 S/2 S]:
    top, bottom = img |> split(b c (2 h) w -> b c h w)
    return top - bottom

(2 h) on an axis of size 16: the checker works out h = 8, makes two results of [b c h w], and checks that two names receive them. In a transformer, split is how one matrix product makes the queries, keys and values: q, k, v = x @ w_qkv |> split(b n (3 H Dh) -> b H n Dh).

Operations along an axis

Many operations work along one axis, and take it by name with over=. The spotter’s class scores are [b C], and most of what you’d do with them is one of these:

dim S = 16
dim C = 4

fn probs(scores: [b C]) -> [b C] Prob over C = softmax(scores, over=C)

fn best(scores: [b C]) -> [b] Index[C] = argmax(scores, over=C)

## The two likeliest shapes of each image.
fn best_two(scores: [b C]) -> [b 2] Index[C]:
    values, ids = topk(scores, 2, over=C)
    return ids

## Two batches of images, one after the other.
fn together(a: [b 1 S S], c: [m 1 S S]) -> [b+m 1 S S] = cat(a, c, over=b)

fn running_total(x: [b n]) -> [b n] = cumsum(x, over=n)

cat adds up the sizes along its axis, so the result is [b+m 1 S S]. stack(a, c, over=k) puts its arguments along a new first axis named k. topk gives two results, the values and their positions.

Indexing and slicing

Indexing removes an axis, and slicing shortens one:

dim S = 16
dim D = 32

fn first_image(img: [b 1 S S]) -> [1 S S] = img[0]

fn first_patch(x: [b n D]) -> [b D] = x[:, 0]

fn drop_last(x: [b n D]) -> [b n-1 D] = x[:, :-1]

fn shifted(x: [b n D]) -> [b n-1 D] = x[:, 1:]

The checker works out a slice’s size when it can prove the slice fits the axis: :-1 and 1: on an axis n both have n-1 elements. A slice that might run past the end, such as x[:, 0:3] on an axis whose size is generic, gets an unknown size.

The spotter so far

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
## 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)

Exercises

1. Change the patches to 8 × 8 pixels. What do you change in spotter.ss, what are N and P*P now, and does patchify need to change?

2. Write unpatchify(x: [b N P*P]) -> Images, which puts the patches back together, so that unpatchify(patchify(img)) is img. Hint: from one axis of size 16, the checker can’t tell how many patches go down and how many across, so give one of them as a keyword. A keyword size can be a number or a dim, but not S / P, which is a float.

3. This first try at the token-mixing MLP uses rearrange and @ instead of einsum. It has a mistake. Find it with the checker, and fix it:

   dim N = 16
   dim D = 32
   dim Dt = 64
   
   fn mix_patches(x: [b N D], w1: [N Dt], w2: [Dt N]) -> [b N D]:
       h = x |> rearrange(b N D -> b D N) |> matmul(w1) |> gelu |> matmul(w2)
       return x + rearrange(h, b N D -> b D N)
   

4. Write ink_map(img: [b 1 S S]), the brightest pixel of each patch, as a 4 × 4 grid per image. What’s its return type?