Calling Python and extern fn

At the end of Chapter 9 the spotter was complete: forward, loss, predict and a way to measure it. But its weights are still whatever you pass in, and nothing has trained them. Python is good at that part: loading data, running an optimizer, saving files. ShapeScript doesn’t replace Python; it lives inside it.

This chapter goes both ways. First Python calls the spotter, to train it. Then the spotter calls Python, to load the weights it was trained with: untyped, with import python, and typed, with extern fn.

Training the spotter from Python

Chapter 7 wrote shapes.py, which draws images of shapes with their labels. Create train.py next to it and spotter.ss:

import torch
from shapes import draw
from spotter import loss, predict, per_class_accuracy, confusion

img, y = draw(2000)                     # training images, as bytes, and labels
test_img, test_y = draw(500, seed=1)    # images the model never trains on

g = torch.Generator().manual_seed(0)
def weights(*shape):
    return (torch.randn(*shape, generator=g) / shape[-2] ** 0.5).requires_grad_()

# In the order the spotter's functions take them: w_e, tw1, tw2, cw1, cw2, w_h.
w = [weights(16, 32), weights(2, 16, 64), weights(2, 64, 16),
     weights(2, 32, 64), weights(2, 64, 32), weights(32, 4)]
opt = torch.optim.Adam(w, lr=3e-3)

for step in range(201):
    batch = torch.randint(0, 2000, (64,), generator=g)
    l = loss(img[batch], y[batch], *w)
    opt.zero_grad()
    l.backward()
    opt.step()
    if step % 50 == 0:
        print(f"step {step}: loss {l.item():.3f}")

with torch.no_grad():
    pred = predict(test_img, *w)
print(f"test accuracy: {(pred == test_y).float().mean().item():.1%}")
print(per_class_accuracy(pred, test_y))
print(confusion(pred, test_y))
torch.save(w, "weights.pt")

Each weight starts as random numbers divided by the square root of the axis it’s multiplied along, so the values that flow through the model start out about the size of the inputs. *w passes the six weights in the order the signatures list them. Run it:

$ shapescript run train.py
step 0: loss 1.606
step 50: loss 0.568
step 100: loss 0.151
step 150: loss 0.083
step 200: loss 0.072
test accuracy: 96.4%
tensor([0.9191, 0.9672, 0.9750, 1.0000])
tensor([[125,   1,  10,   0],
        [  0, 118,   1,   3],
        [  2,   1, 117,   0],
        [  0,   0,   0, 122]])

It takes a few seconds on a laptop. The confusion matrix says where the 3.6% of mistakes are: mostly squares taken for pluses, 10 of them, with a few other slips scattered round. At the smallest size, a square outline and a plus are both a few bright pixels in a 4×4 patch, and the model hasn’t quite learned to tell them apart.

Nothing here is special to ShapeScript: loss is an ordinary Python function that calls torch, so autograd, Adam and torch.save work as usual. The last line saves the six trained weights to weights.pt.

import python

Now the other way round. The weights are in a file, and the spotter could load them itself. import python NAME makes a Python module available:

import python torch

## The likeliest class of each image, with the weights `train.py` saved.
fn predict_saved(img: Images u8) -> [b] Index[C]:
    w = torch.load("weights.pt")
    return predict(img, w[0] as [P*P D], w[1] as [L N Dt], w[2] as [L Dt N],
                   w[3] as [L D Dc], w[4] as [L Dc D], w[5] as [D C])

torch.load(...) is a Python value: the checker knows nothing about it, not even whether it’s a tensor. Indexing it, w[0], or doing arithmetic on it gives another Python value, and you can pass Python values back to Python freely. But where a shape matters, such as passing it to a parameter with a shape, returning it as a tensor, or reducing it over= an axis, it’s an error until an as gives it a type:

import python torch

dim D = 32
dim C = 4

fn saved_head() -> [D C]:
    return torch.load("weights.pt")[5]

$ shapescript check untyped.ss
error[E0413]: the shape of this Python value is unknown
 --> untyped.ss:7:12
  |
7 |     return torch.load("weights.pt")[5]
  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^ the return type is `[D C]`
  |
  = help: assert its shape with `as [D C]`

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

After the as, everything is checked statically again: predict gets six typed weights, and the checker follows them through the model as before. That’s the pattern for the boundary with Python: one as where a value comes in, and types from there on.

Each as is also checked when the program runs (Chapter 7). Suppose you widen the tokens to dim D = 48, but forget to retrain, so weights.pt still holds the old weights:

ShapeError: spotter.ss:86:25: expected [P*P D] = (16, 48), got shape (16, 32)

The error is at the first weight that doesn’t fit, on the line that loads it, before any of the model runs.

Python values as host values

Python values also work where the program needs a host value: as an if condition, the end of a loop, or what a for loops over. shapes.py has a list of the class names, which a function can print:

import python shapes

fn show_classes():
    for name in shapes.CLASSES:
        print(name)

import python shapes finds shapes.py as Python’s import would, on the Python path, which includes the folder of the script you run. Here name is a Python value too, and print takes anything.

extern fn: typed Python functions

import python is fine for one-off calls. For a function you call often, give it a type with extern fn, naming the Python function in quotes:

## The complementary error function, $1 - \operatorname{erf}(x)$.
extern fn erfc(x: Real) -> NonNeg = "torch.special.erfc"

dim D = 64

fn tail(x: [b D]) -> [b D] NonNeg = erfc(x)

Calls are checked against the signature like any other, and compile to a call of the Python function, with its module imported. The signature is trusted: what the function returns isn’t checked when it runs, so get it right.

Many built-ins are exactly this. exp, log, sigmoid, gelu and the like are extern fn declarations in the prelude, a file every program sees, and linear, layer_norm, rms_norm, conv2d, embedding, scaled_dot_product_attention and dropout are in a second one for torch.nn.functional. The gelu in mix_patches is one of them: in the editor, go to definition on gelu opens its declaration. A declaration in your file with the same name hides the prelude’s.

What an extern fn signature can say

Since there’s no body to check, an extern fn’s signature can say things a fn’s can’t. The spotter has a use for two of them.

The first is a test of how robust the spotter is. A shape moved one pixel to the right is still the same shape, and a good spotter should say so. torch.roll moves the elements of a tensor along one axis, wrapping the last ones round to the front:

## Rolls the axis `dims` round by `shifts` places.
extern fn roll(input: [...] u8, shifts: Size[K], dims: Axis[n]) -> [...] u8 = "torch.roll"

## Moves each image one pixel to the right; the last column wraps round.
fn nudge(img: [b c h w] u8) -> [b c h w] u8 = roll(img, 1, dims=w)

• ... stands for any number of other axes, the same in every type of the signature. roll takes any shape and gives back the same. The dtype is part of the type, as anywhere: this roll is for the spotter’s bytes.
• dims: Axis[n] takes an axis name, dims=w, and the Python function gets its position, here dims=3. That’s how torch’s positional dim= arguments become named ones. n is the size of that axis, a new name of the signature’s own. As anywhere in a signature, a dim’s name isn’t new: in spotter.ss, Axis[N] would mean an axis of size N = 16, and dims=w would be an error, whose help says so.
• shifts: Size[K] takes a whole number known before the program runs, which becomes the size K in the types.

Why does nudge take [b c h w] and not Images? Because Images is [b 1 S S], which has two axes named S, and dims=S couldn’t say which. In nudge, the axes have names of their own, and a call with Images matches h and w with the two Ses by position, as Chapter 4 described.

Some torch functions take a list of axes, never one number. torch.flip is one: torch.flip(x, 3) is a TypeError. For them, write Axes[n], which takes one axis name too, and passes a list of its position, dims=[3]:

extern fn flip(input: [...] u8, dims: Axes[n]) -> [...] u8 = "torch.flip"

## Each image, mirrored left to right.
fn mirror(img: [b c h w] u8) -> [b c h w] u8 = flip(img, dims=w)

Whether a function wants Axis or Axes is in its documentation, not its name, and the checker can’t know: the signature is trusted. Get it wrong, and the call fails inside torch. That’s the price of an extern fn, and why it’s worth a short test.

The second is a small preview of an image, made by pooling: each output pixel is the brightest pixel of a 3×3 square, with squares every 2 pixels.

## The largest value of each `kernel_size`-wide square, every `stride` pixels.
extern fn max_pool2d(input: [b c h w], kernel_size: Size[K], stride: Size[St]) -> [b c (h - K)/St + 1 (w - K)/St + 1] = "torch.nn.functional.max_pool2d"

## A 7×7 summary of each image.
fn thumbnail(img: Images) -> [b 1 7 7] = max_pool2d(img, 3, stride=2)

Division in the return type rounds down, as in torch’s formulas for convolutions and pooling. For a 16-pixel side, (16 − 3) / 2 + 1 is 6.5 + 1 rounded down: 7. The checker works that out from the signature, so -> [b 1 7 7] checks, and [b 1 8 8] wouldn’t.

Here’s a script that loads the trained spotter, tries it on nudged images, and makes thumbnails:

import torch
from shapes import draw
from spotter import predict_saved, nudge, thumbnail

img, y = draw(500, seed=1)
pred = predict_saved(img)
print(f"accuracy: {(pred == y).float().mean().item():.1%}")
print(f"nudged: {(predict_saved(nudge(img)) == y).float().mean().item():.1%}")
print(thumbnail(img.float()).shape)

$ shapescript run check.py
accuracy: 96.4%
nudged: 93.0%
torch.Size([500, 1, 7, 7])

The same accuracy as at the end of training, so the weights came back intact. Moved by one pixel, the images lose 3.4 points: the spotter has learned a little about where the shapes were in its training images, not only what they look like. An exercise below makes it more robust.

Python calling ShapeScript, again

train.py imported spotter.ss as if it were a Python module. A few more details are worth knowing:

• Every ShapeScript function checks its arguments when it’s called, and on the declared device: a tensor on the CPU passed to a program declared devices cuda[0] raises shapescript.DeviceError naming the argument, instead of failing somewhere deep inside.
• There’s no default device. A Python program declares it once, with shapescript.devices("cuda[0]"), unless a .ss file it imports has a devices line, as spotter.ss does.
• shapescript run sets up the import hook for you. A script run with plain python calls shapescript.install() before its first ShapeScript import.
• shapescript.compile(src, path) returns the Python and the diagnostics of some ShapeScript, without running anything.

The spotter so far

devices cpu

import python torch

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

## The likeliest class of each image, with the weights `train.py` saved.
fn predict_saved(img: Images u8) -> [b] Index[C]:
    w = torch.load("weights.pt")
    return predict(img, w[0] as [P*P D], w[1] as [L N Dt], w[2] as [L Dt N],
                   w[3] as [L D Dc], w[4] as [L Dc D], w[5] as [D C])

## Rolls the axis `dims` round by `shifts` places.
extern fn roll(input: [...] u8, shifts: Size[K], dims: Axis[n]) -> [...] u8 = "torch.roll"

## Moves each image one pixel to the right; the last column wraps round.
fn nudge(img: [b c h w] u8) -> [b c h w] u8 = roll(img, 1, dims=w)

## The largest value of each `kernel_size`-wide square, every `stride` pixels.
extern fn max_pool2d(input: [b c h w], kernel_size: Size[K], stride: Size[St]) -> [b c (h - K)/St + 1 (w - K)/St + 1] = "torch.nn.functional.max_pool2d"

## A 7×7 summary of each image.
fn thumbnail(img: Images) -> [b 1 7 7] = max_pool2d(img, 3, stride=2)

Exercises

1. Make the spotter’s predictions steadier: write predict_both, which takes the same arguments as predict, adds the class scores of each image to those of the image nudged one pixel, and returns the likeliest class of the sum.

2. Write an extern fn for torch.nn.functional.avg_pool2d, which takes the same arguments as max_pool2d but averages each square, and use it for tiny(img: Images) -> [b 1 4 4], which averages each 4×4 patch. What would tiny’s return type be with a kernel of 5 and a stride of 3?

3. A colleague saved 100 test images with torch.save(img, "test.pt"). Fix the error, keeping the signature, and say what happens at run time if the file holds 28×28 images:

   import python torch
   
   dim S = 16
   
   fn test_images() -> [?k 1 S S] u8:
       return torch.load("test.pt")