Control Flow and Iterator Chains

The spotter works, but it’s shallow: one block that mixes the patches, one that mixes the channels. Real mixers stack several. And when it gets an image wrong, a single accuracy figure doesn’t say which shapes it confuses. This chapter fixes both. Stacking blocks needs a for loop, whose steps run one after another. Measuring each class needs an iterator chain, which computes something for every element along an axis, all at once. Along the way we’ll meet if.

Stacking blocks with a loop

for i in a..b: runs its block for each whole number from a up to, but not including, b, like Python’s range(a, b).

To stack L blocks, each block needs its own weights. Rather than eight weight parameters for two blocks, and more for three, give each block weight a leading axis L: tw1: [L N Dt] holds the patch-mixing weights of every block, and tw1[l], indexing that axis, is block l’s [N Dt]. Here is the new forward, which replaces Chapter 8’s. Add L to the dims, and change forward, loss and predict in spotter.ss:

## Number of mixer blocks.
dim L = 2

## The class scores of each image.
fn forward(img: Images u8, w_e: [P*P D], tw1: [L N Dt], tw2: [L Dt N],
           cw1: [L D Dc], cw2: [L Dc D], w_h: [D C]) -> [b C] Logit over C:
    x = embed(img |> to Unit |> to Signed, w_e)
    for l in 0..L:
        x = x |> mix_patches(tw1[l], tw2[l]) |> mix_channels(cw1[l], cw2[l])
    return x |> norm |> head(w_h)

forward is no longer one pipeline: it embeds the patches, runs the blocks in turn, then normalizes and classifies. The ends of a range are host numbers: ints, dims or sizes, never tensors, since a loop counts on the host. l is an int, and tw1[l] is a [N Dt], which is what mix_patches takes. There’s no while, break or continue; a return inside a loop returns from the function.

loss and predict take the stacked weights too, and pass them on:

## The mean cross-entropy of the scores against the labels.
fn loss(img: Images u8, y: [b] Index[C], w_e: [P*P D], tw1: [L N Dt], tw2: [L Dt N],
        cw1: [L D Dc], cw2: [L Dc D], w_h: [D C]) -> [] NonNeg:
    return cross_entropy(forward(img, w_e, tw1, tw2, cw1, cw2, w_h), y, over=C)

## The likeliest class of each image.
fn predict(img: Images u8, w_e: [P*P D], tw1: [L N Dt], tw2: [L Dt N],
           cw1: [L D Dc], cw2: [L Dc D], w_h: [D C]) -> [b] Index[C]:
    return argmax(forward(img, w_e, tw1, tw2, cw1, cw2, w_h), over=C)

A deeper spotter is now one line: dim L = 4. The checker then expects every block weight with 4 in front, and so does the runtime check on the way in from Python.

What a loop may change

A loop’s block may give a new value to a variable from outside it, as x = x |> ... does. The new value must have the same shape, axis names and dtype, so that the code after the loop, and the next time round, work however many times the loop ran. Suppose a stray |> pool, which belongs after the loop, ends up inside it:

fn blocks(x: [b N D], tw1: [L N Dt], tw2: [L Dt N]) -> [b N D]:
    for l in 0..L:
        x = x |> mix_patches(tw1[l], tw2[l]) |> pool
    return x

$ shapescript check loop_bad.ss
error[E0801]: `x` can't change type in a loop
 --> loop_bad.ss:3:13
  |
1 | fn blocks(x: [b N D], tw1: [L N Dt], tw2: [L Dt N]) -> [b N D]:
  |           - `x` is `[b N D]` here
2 |     for l in 0..L:
3 |         x = x |> mix_patches(tw1[l], tw2[l]) |> pool
  |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ this is `[b D]`, but `x` is `[b N D]`
  |
  = help: a value from outside a loop must have the same type on every iteration; give this one a new name

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

The first time round, x would become a [b D], and the second block can’t take that. To compute something of another type, give it a new name. Names first bound inside a block are gone after it.

Loops also run over a dataset’s batches, for batch in data |> batches(64):, which Chapter 12 covers, and over Python values (Chapter 10).

Conditions

if takes a bool or a [] mask, and else: goes on its own line. There’s no elif; nest an if in the else block instead. Say you want to try other activations in the channel MLP, chosen by a number from Python:

dim Dc = 64

fn activation(h: [b n Dc], kind: int) -> [b n Dc]:
    if kind == 0:
        return gelu(h)
    else:
        if kind == 1:
            return relu(h)
        else:
            return tanh(h)

## Shrinks `h` if any of it is large.
fn tame(h: [b n Dc]) -> [b n Dc]:
    if max(abs(h)) > 10:
        h = h * 0.1
    return h

kind == 0 compares two host numbers, a bool. max(abs(h)) > 10 is a [] mask, a value on the device: reading it to decide which way to go makes the program wait for the device, as if tensor.item() > 10 would in PyTorch. A mask with axes isn’t true or false; reduce it with any or all first.

Measuring the spotter with iterator chains

Comparing predictions with labels gives the overall accuracy: mean(pred == y to f32). What we want is one number per class: of the images that really are squares, how many did the spotter call squares? In PyTorch you’d write a loop over the classes, or some careful indexing. In ShapeScript it’s an iterator chain: you write the computation for one class, and say what to iterate.

## The fraction of each class's images that were predicted right.
fn per_class_accuracy(pred: [b] Index[C], y: [b] Index[C]) -> [C]:
    return C |> map(c => (sum(pred == c and y == c) to f32) / (sum(y == c) to f32))

C |> map(c => e) builds a tensor whose first axis is C, computing e for every index c along it. c => e is a lambda: a parameter, =>, and what to compute from it. Inside, c is an Index[C], a class, so pred == c is a mask of the images predicted as that class. The result has the iterated axis in front, still named C.

The parentheses matter: to binds more tightly than == and and, so sum(pred == c and y == c) to f32 without them would convert c, not the count.

The checker has a warning:

$ shapescript check accuracy.ss
warning[W0501]: `sum(y == c) to f32` can be exactly 0
 --> accuracy.ss:4:63
  |
4 |     return C |> map(c => (sum(pred == c and y == c) to f32) / (sum(y == c) to f32))
  |                                                               ^^^^^^^^^^^^^^^^^^^^ dividing by 0 gives inf or NaN
  |
  = note: it's `NonNeg`, so its values are in [0, inf)
  = help: if it's never 0, say so with `as Pos`
  = note: in stage 1 of this pipeline, whose input is the indices of `C`

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

It’s right. A batch of 64 random shapes can easily have no crosses in it, and then this class’s accuracy is 0 / 0, a NaN, which would spoil any average you took of it. Divide by at least 1, and an absent class counts as 0:

## The fraction of each class's images that were predicted right.
fn per_class_accuracy(pred: [b] Index[C], y: [b] Index[C]) -> [C]:
    return C |> map(c => (sum(pred == c and y == c) to f32) / maximum(sum(y == c) to f32, 1.0))

A chain can go inside a chain. The confusion matrix counts, for each true class and each predicted class, how many images had that pair. Its diagonal holds the right answers, and anything off it is a mistake, by kind:

## How often each true class (rows) was predicted as each class (columns).
fn confusion(pred: [b] Index[C], y: [b] Index[C]) -> [C C] i64:
    return C |> map(t => C |> map(p => sum(y == t and pred == p)))

A sum of a mask counts its true elements, so the counts are whole numbers, i64. Chapter 10 trains the spotter and prints both; the confusion matrix shows that small squares and small pluses are the shapes it mixes up.

Stages

A chain can start with a size (C, n, 3), a range (0..n) or a tensor, and has four stages:

• map(f): f of every element, stacked along the iterated axis.
• filter(f): the elements for which f is true. How many there are is only known when the program runs, so their axis is an unknown size, ?k, as Chapter 3 described.
• zip(other): pairs of elements from two iterations of the same length, for a lambda with two parameters, (a, b) => ….
• enumerate: pairs of each element’s index and the element.

Starting from a tensor, over= says which axis to iterate, and the elements are the slices along it:

dim C = 4
dim D = 32

## The true class of each image the spotter got wrong.
fn missed(pred: [b] Index[C], y: [b] Index[C]) -> [?k] Index[C]:
    return y
        |> zip(pred, over=b)
        |> filter((t, p) => t != p)
        |> map((t, p) => t)

## Each token's channels, scaled by the token's position.
fn by_position(x: [n D]) -> [n D] = x |> enumerate(over=n) |> map((i, row) => row * (i to f32))

fn smoothed(x: [n D]) -> [n D] = x |> map(gelu, over=n)

The last one gives map a function by name instead of a lambda, as a pipe stage would be: map(gelu) is map(v => gelu(v)), and map(to f32) is map(v => v to f32). In the editor, each line of a chain shows its shape at the end, as for any pipe.

How chains run

A lambda that’s elementwise, like by_position’s, becomes ordinary tensor operations on whole tensors: here one multiplication of x by the positions, as a column. Anything else, like the sums in confusion, becomes torch.vmap, which autograd and torch.export both understand. Either way, no Python loop runs: the 16 entries of the confusion matrix are computed at once, not one after another.

Elements can’t depend on each other

Because the elements are computed at once, one can’t read another. A running sum is a recurrence, which is a loop’s job:

fn running(x: [n]) -> [n]:
    y = n |> map(i => y[i - 1] + x[i])
    return y

$ shapescript check recurrence.ss
error[E2402]: a lambda of an iterator stage can't read `y`, the result it builds
 --> recurrence.ss:2:23
  |
2 |     y = n |> map(i => y[i - 1] + x[i])
  |                       ^ the result of other elements
  |
  = note: an iterator stage computes every element at once, so no element can depend on another
  = help: to compute each element from the one before, use a `for` loop, which runs in order
  = note: in stage 1 of this pipeline, whose input is the indices of `n`

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

(For this one, there’s a built-in: cumsum(x, over=n).) For the same reason a lambda can’t print, call Python, draw random numbers or change a model.

Loops that should be iterator chains

The other way round, the checker notices a loop whose steps don’t depend on each other, and suggests a chain. Here is pool written as a loop over the patches:

dim D = 32

fn pool_loop(x: [b n D]) -> [b D]:
    acc = x[:, 0] * 0
    for i in 0..n:
        acc = acc + x[:, i]
    return acc / (n as Pos)

$ shapescript check hint.ss
hint[H2401]: the iterations of this loop over `n` don't depend on each other
 --> hint.ss:5:5
  |
5 |     for i in 0..n:
  |     ^^^^^^^^^^^^^ runs one iteration after another
6 |         acc = acc + x[:, i]
  |                     ------- the term added to `acc`
  |
  = note: each iteration adds a term to `acc` that doesn't read what the others added, so every term can be computed at once
  = help: compute every term at once with `map`: `acc = acc + (n |> map(i => x[:, i]) |> sum(over=n))`

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

The mixer blocks are different: block 2 takes block 1’s output, so they really must run in order, and the checker says nothing about that loop.

Host scalars

int, float and bool are host scalars: plain Python values for counters, rates and flags, which never go to the GPU. Numbers in the source, loop variables such as l, and sizes used as values are host scalars. Arithmetic on them gives host scalars, and arithmetic with a tensor gives a tensor.

A tensor never turns into a host scalar by itself, because that would make the program wait for the device. Running one more block per mistake, say, is an error:

dim N = 16
dim D = 32
dim C = 4

fn repeat_n(x: [b N D], n: int) -> [b N D]:
    for i in 0..n:
        x = x * 0.5
    return x

fn bad(x: [b N D], pred: [b] Index[C], y: [b] Index[C]) -> [b N D] = repeat_n(x, sum(pred != y))

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
## Hidden width of the MLP that mixes channels.
dim Dc = 64
## Number of classes: square, disk, plus and cross.
dim C = 4
## Number of mixer blocks.
dim L = 2

## 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] Logit over 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 u8, w_e: [P*P D], tw1: [L N Dt], tw2: [L Dt N],
           cw1: [L D Dc], cw2: [L Dc D], w_h: [D C]) -> [b C] Logit over C:
    x = embed(img |> to Unit |> to Signed, w_e)
    for l in 0..L:
        x = x |> mix_patches(tw1[l], tw2[l]) |> mix_channels(cw1[l], cw2[l])
    return x |> norm |> head(w_h)

## The mean cross-entropy of the scores against the labels.
fn loss(img: Images u8, y: [b] Index[C], w_e: [P*P D], tw1: [L N Dt], tw2: [L Dt N],
        cw1: [L D Dc], cw2: [L Dc D], w_h: [D C]) -> [] NonNeg:
    return cross_entropy(forward(img, w_e, tw1, tw2, cw1, cw2, w_h), y, over=C)

## The likeliest class of each image.
fn predict(img: Images u8, w_e: [P*P D], tw1: [L N Dt], tw2: [L Dt N],
           cw1: [L D Dc], cw2: [L Dc D], w_h: [D C]) -> [b] Index[C]:
    return argmax(forward(img, w_e, tw1, tw2, cw1, cw2, w_h), over=C)

## The fraction of each class's images that were predicted right.
fn per_class_accuracy(pred: [b] Index[C], y: [b] Index[C]) -> [C]:
    return C |> map(c => (sum(pred == c and y == c) to f32) / maximum(sum(y == c) to f32, 1.0))

## How often each true class (rows) was predicted as each class (columns).
fn confusion(pred: [b] Index[C], y: [b] Index[C]) -> [C C] i64:
    return C |> map(t => C |> map(p => sum(y == t and pred == p)))

Exercises

1. Add accuracy(pred: [b] Index[C], y: [b] Index[C]) -> [] to the spotter: the fraction of all images predicted right.

2. Write class_counts(y: [b] Index[C]) -> [C] i64, how many images of each class a batch has, with map. Then get the same numbers from confusion instead, without a chain.

3. A smaller spotter could use one set of block weights for every block. Write forward_shared, which takes tw1: [N Dt], tw2: [Dt N], cw1: [D Dc] and cw2: [Dc D], and applies the same block L times.

4. This function, meant to drop the first patch after each block, doesn’t check. Why?

   dim D = 32
   dim L = 2
   
   fn shrink(x: [b n D]) -> [b n D]:
       for l in 0..L:
           x = x[:, 1:]
       return x