A First Program: a Classifier

Let’s write something real: a small neural network that sorts points in the plane into three classes. We’ll write the model’s maths in ShapeScript and drive it from a short Python script, so you can see how the two fit together. Many of the ideas here get a chapter of their own later; this one is a tour.

Setting up

Make a folder called classifier, and in it a file classifier.ss. Start with the sizes:

devices cpu

## Number of input features.
dim D = 2
## Width of the hidden layer.
dim Hd = 32
## Number of classes.
dim C = 3

A dim is a size with a name. Types use these names, so the program says which 32 it means, not just 32.

The forward pass

The network takes a batch of points, multiplies by a weight matrix, applies a ReLU, and multiplies by another:

## The class scores of each example.
fn forward(x: [b D], w1: [D Hd], w2: [Hd C]) -> [b C] Logit over C:
    return x @ w1 |> relu |> matmul(w2)

Read the signature first:

• x: [b D] is a batch of points. b is a name that isn’t declared anywhere, so it’s a generic axis: the caller decides its size, and every call can use a different one.
• w1: [D Hd] and w2: [Hd C] are the weights, with fixed sizes.
• -> [b C] Logit over C is the result: one row of scores per point, one score per class. Logit over C says what the numbers mean: unnormalized scores over the axis C, the kind you’d pass to softmax or cross_entropy.

Then the body. x @ w1 is a matrix product, as in PyTorch: [b D] @ [D Hd] is [b Hd]. |> relu passes that to relu, and |> matmul(w2) passes the result as the first argument of matmul, so the line means matmul(relu(x @ w1), w2), which reads in the order it happens.

The checker follows every shape: [b D] @ [D Hd] is [b Hd], relu keeps it, and [b Hd] @ [Hd C] is [b C], which is what the signature promised.

The loss and predictions

## The mean cross-entropy of the scores against the labels.
fn loss(x: [b D], y: [b] Index[C], w1: [D Hd], w2: [Hd C]) -> [] NonNeg:
    return cross_entropy(forward(x, w1, w2), y, over=C)

## The likeliest class of each example.
fn predict(x: [b D], w1: [D Hd], w2: [Hd C]) -> [b] Index[C]:
    return argmax(forward(x, w1, w2), over=C)

y: [b] Index[C] is one label per point, and Index[C] says each is a whole number from 0 to C - 1, a position along the axis C. -> [] NonNeg is a scalar, a tensor with no axes, that is never negative.

cross_entropy(..., over=C) and argmax(..., over=C) take the class axis by name. There’s no dim=1 to get wrong: if the axis isn’t called C, it’s an error.

Check the file:

shapescript check classifier.ss

No output: it’s correct.

Training it from Python

Create train.py next to it:

import torch
from classifier import loss, predict

# Three blobs of points in the plane, one per class.
g = torch.Generator().manual_seed(0)
centers = torch.tensor([[0.0, 2.0], [-2.0, -1.0], [2.0, -1.0]])
y = torch.randint(0, 3, (300,), generator=g)
x = centers[y] + 0.7 * torch.randn(300, 2, generator=g)

w1 = (0.5 * torch.randn(2, 32, generator=g)).requires_grad_()
w2 = (0.5 * torch.randn(32, 3, generator=g)).requires_grad_()

for step in range(201):
    l = loss(x, y, w1, w2)
    l.backward()
    with torch.no_grad():
        for w in (w1, w2):
            w -= 0.1 * w.grad
            w.grad = None
    if step % 50 == 0:
        print(f"step {step}: loss {l.item():.3f}")

accuracy = (predict(x, w1, w2) == y).float().mean().item()
print(f"accuracy: {accuracy:.1%}")

and run it:

$ shapescript run train.py
step 0: loss 2.059
step 50: loss 0.094
step 100: loss 0.081
step 150: loss 0.077
step 200: loss 0.074
accuracy: 97.3%

from classifier import loss, predict compiled classifier.ss to Python, and the functions it imported are ordinary PyTorch code, so autograd, .backward() and everything else work as usual. (Chapter 12 writes this loop in ShapeScript instead, with optimizers, datasets and checkpoints.)

What ShapeScript adds at run time is a check on the way in. Swap the weights by mistake, loss(x, y, w2, w1), and the call fails at once, raising a ShapeError that names the parameter, its type and both shapes. shapescript run shows it as the checker shows its errors, with the argument underlined and the parameter it was given for:

$ shapescript run train.py
error: argument `w1` of `loss` should be [D Hd] = (2, 32), got shape (32, 3)
  --> train.py:14:20
   |
14 |     l = loss(x, y, w2, w1)
   |                    ^^ expected [D Hd] = (2, 32), got shape (32, 3)
   |
  ::: classifier.ss:15:36
   |
15 | fn loss(x: [b D], y: [b] Index[C], w1: [D Hd], w2: [Hd C]) -> [] NonNeg:
   |                                    ---------- parameter declared here
   |
   = note: run with `SHAPESCRIPT_TRACEBACK=1` to see the Python traceback

Making mistakes on purpose

The real point of the checker is the mistakes it catches before anything runs. Here are three that happen in real code.

Applying softmax before the loss. cross_entropy takes logits, and applies log_softmax itself. Passing it probabilities still trains, badly, and nothing in PyTorch complains. ShapeScript knows what softmax returns:

dim D = 2
dim Hd = 32
dim C = 3

fn forward(x: [b D], w1: [D Hd], w2: [Hd C]) -> [b C] Logit over C:
    return x @ w1 |> relu |> matmul(w2)

fn loss(x: [b D], y: [b] Index[C], w1: [D Hd], w2: [Hd C]) -> [] NonNeg:
    p = softmax(forward(x, w1, w2), over=C)
    return cross_entropy(p, y, over=C)

$ shapescript check softmax_twice.ss
error[E0501]: expected `Logit over C`, got `Prob over C`
  --> softmax_twice.ss:10:26
   |
10 |     return cross_entropy(p, y, over=C)
   |                          ^ this is `[b C] Prob over C`
   |
   = help: softmax applied before cross_entropy?

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

Mixing up axes. Suppose you average the scores over the batch to get a “typical” score per class, but write the wrong axis:

dim C = 3

fn mean_scores(z: [b C]) -> [C]:
    return mean(z, over=C)

$ shapescript check wrong_axis.ss
error[E0401]: mismatched shape: expected `[C]`, found `[b]`
 --> wrong_axis.ss:4:12
  |
4 |     return mean(z, over=C)
  |            ^^^^^^^^^^^^^^^ the return type is `[C]`
  |
  = note: axis 1: `b` is not known to equal `C` = 3

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

Averaging over C leaves [b], which isn’t [C]. The fix is over=b.

A label from the wrong set. Labels typed Index[C] are positions along an axis of size C. Using them to index an axis of another size is an error, where PyTorch would read the wrong rows or fail with an index error at run time:

dim C = 3
dim K = 5

fn pick(table: [K 4], y: [b] Index[C]) -> [b 4]:
    return table[y]

$ shapescript check wrong_index.ss
error[E0507]: expected `Index[K]`, got `Index[C]`
 --> wrong_index.ss:5:18
  |
5 |     return table[y]
  |            ----- ^ this is `[b] Index[C]`
  |            |
  |            the axis it indexes is `K` = 5
  |
  = note: the sizes differ: `C` = 3, not `K` = 5

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

Exercises

1. Add a bias to the hidden layer: a parameter b1: [Hd] added after x @ w1. Where does the bias go in the signature, and does + need an unsqueeze?
2. Write accuracy(x: [b D], y: [b] Index[C], w1: [D Hd], w2: [Hd C]), returning the fraction of correct predictions as a []. Hint: comparing two tensors with == gives a mask, and to f32 turns a mask into numbers you can average.
3. Predict the type of forward(x, w1, w2) where x: [16 D].