Assertions with as and Runtime Checks

The spotter has a whole forward pass now: patches, a token per patch, a mixer block, and class scores. Inside spotter.ss, the checker has proved every shape. But the images and weights come from Python, and the checker can’t see Python. This chapter runs the spotter from Python for the first time and covers what protects it there: the checks that run when the program does, and as, which tells the checker something it can’t prove.

Running the spotter

First we need images. Make shapes.py next to spotter.ss. It draws squares, disks, pluses and crosses of random sizes in random places, on a noisy background, and returns them as bytes, the way images are usually stored:

import torch

CLASSES = ["square", "disk", "plus", "cross"]


def draw(n, seed=0):
    """`n` 16×16 grayscale images of shapes, as bytes, and their labels."""
    g = torch.Generator().manual_seed(seed)
    y = torch.randint(0, 4, (n,), generator=g)
    r = torch.randint(2, 7, (n, 1, 1), generator=g).float()  # half the size
    cy = torch.randint(5, 11, (n, 1, 1), generator=g).float()
    cx = torch.randint(5, 11, (n, 1, 1), generator=g).float()
    i = torch.arange(16).float()
    dy, dx = i.view(1, 16, 1) - cy, i.view(1, 1, 16) - cx
    shapes = torch.stack([
        (torch.maximum(dy.abs(), dx.abs()) - r).abs() < 0.6,         # square
        dy**2 + dx**2 <= r**2,                                        # disk
        ((dy.abs() < 1) | (dx.abs() < 1)) & (dy.abs() <= r) & (dx.abs() <= r),  # plus
        ((dy - dx).abs() < 1) | ((dy + dx).abs() < 1),                # cross
    ])
    ink = shapes[y, torch.arange(n)]
    ink[y == 3] &= (dy.abs() <= r)[y == 3]
    pixels = ink.float() * 140 + 110 * torch.rand(n, 16, 16, generator=g)
    return pixels.to(torch.uint8).unsqueeze(1), y

There’s nothing to learn from that file except what it returns: a (n, 1, 16, 16) tensor of bytes, and n labels from 0 to 3, one per class in CLASSES.

Now run.py, which pushes four images through the spotter with random weights:

import torch
from shapes import CLASSES, draw
from spotter import forward

img, y = draw(4)
img = img / 255  # bytes to numbers in [0, 1]

g = torch.Generator().manual_seed(0)
w = [torch.randn(16, 32, generator=g), torch.randn(16, 64, generator=g),
     torch.randn(64, 16, generator=g), torch.randn(32, 64, generator=g),
     torch.randn(64, 32, generator=g), torch.randn(32, 4, generator=g)]

scores = forward(img, *w)
print(scores.shape)
for s, label in zip(scores.argmax(-1), y):
    print("guessed", CLASSES[s], "for a", CLASSES[label])

img / 255 turns the bytes into numbers in [0, 1], since forward takes f32 images; Chapter 8 lets the spotter take the bytes itself. The weights are in the order of forward’s parameters: w_e, tw1, tw2, cw1, cw2 and w_h. Run it:

$ shapescript run run.py
torch.Size([4, 4])
guessed disk for a square
guessed disk for a cross
guessed disk for a disk
guessed disk for a square

The spotter runs: four images in, four rows of four scores out. The guesses are no better than chance, because the weights are random: this one happens to think everything is a disk. Chapter 10 trains them.

Runtime checks

forward is ordinary Python now, and Python will call it with anything. So every ShapeScript function checks its arguments when it’s called, whether from Python or from other ShapeScript: the number of axes, each fixed size, that the generic sizes agree across the arguments, and that each holds the right kind of values: floats, integers or booleans, as its dtype says. Here are two mistakes worth making on purpose. A batch of 28×28 images, the size of MNIST’s digits:

forward(torch.rand(4, 1, 28, 28), *w)

and the patch-mixing weights in the wrong order:

w_e, tw1, tw2, cw1, cw2, w_h = w
forward(img, w_e, tw2, tw1, cw1, cw2, w_h)

Each fails at once, at the call, naming the parameter, its type, and both shapes:

ShapeError: spotter.ss:54:12: argument `img` of `forward` should be [b 1 S S] = (?, 1, 16, 16), got shape (4, 1, 28, 28)
ShapeError: spotter.ss:54:39: argument `tw1` of `forward` should be [N Dt] = (16, 64), got shape (64, 16)

(The messages give the file’s full path; it’s shortened here.) Without the checks, both mistakes would still fail, but deep inside: in rearrange, with a message about einops’ pattern, and in an einsum, with one about operand sizes. And if Dt were 16, so that both weights were [16 16], the swap wouldn’t fail at all.

ShapeError is a ValueError, from the shapescript package, as are the other errors in this chapter. A traceback shows the lines of the .ss file, not of generated Python.

as claims a type

The checks on the way in protect a function’s signature. Sometimes you want to go further, and make a claim the checker can’t prove. x as T claims that x already has the type T. The checker believes it from then on, and the compiled program checks it when it runs.

Say the images come from a folder of files, and Python loads them as [b 1 h w], whatever their size. They’re supposed to be 16×16. Add to spotter.ss:

## Images of any size from Python, which must be S × S.
fn from_python(img: [b 1 h w]) -> Images = img as Images

h and w are generic, so they might be 16 at run time, and the claim is allowed. Without the as, returning img would be an error, because h isn’t known to equal S. With it, the rest of the program sees Images, and passes it to forward without complaint.

Unknown sizes

Chapter 3’s nonblank keeps the images that have any ink, so it returns [?k 1 S S]: how many it keeps is only known when it runs. Often you know more. If every image in a batch should have ink, the number kept is the number of images, and you can say so:

## The images, which must all have ink.
fn inked(img: Images) -> Images = nonblank(img) as [b 1 S S]

?k becomes b, the batch size of img. If an image is blank after all, the claim is false, and the program stops there instead of carrying on with three images and four labels.

as is the only way to make a claim the checker can’t prove, and it always compiles to a runtime check. So every place a program relies on an assumption is a visible as, and you can find them all with a search.

What as can claim

An as must be possible. A claim that contradicts what the checker knows is an error, because it can only fail:

dim S = 16
dim Images = [b 1 S S]

fn upscale(img: Images) -> [b 1 28 28]:
    return img as [b 1 28 28]

$ shapescript check upscale.ss
error[E0410]: `[b 1 S S]` can never be `[b 1 28 28]`
 --> upscale.ss:5:12
  |
5 |     return img as [b 1 28 28]
  |            ---^^^^^^^^^^^^^^^
  |            |
  |            this is `[b 1 S S]`
  |
  = note: axis 3: `S` = 16, not 28

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

as never changes a value, so it can’t resize an image. Symbolic sizes are fine, since they might be equal at run time, as h and S were in from_python. That makes as the way to say two generic axes are the same. Blending two batches of images, pixel by pixel, needs them to be equally long:

dim S = 16
dim Images = [b 1 S S]

## The average of two batches of images, which must be equally long.
fn blend(a: Images, c: [m 1 S S]) -> Images = (a + (c as [b 1 S S])) / 2

Without the as, a + c is error E0404 (Chapter 6), because b and m are different axes. With it, the program checks at run time that c has b images.

as also gives a type to a value the checker knows nothing about, such as one from Python (Chapter 10). The parts of the type you don’t write are kept: x as [b C] keeps x’s domain, and x as Unit keeps its shape.

Claiming values

A type can also say what values a tensor holds: its domain, which Chapter 8 is about. Unit means values in [0, 1], which is what forward expects of its pixels. run.py divided by 255 to get there; forgetting to is a classic bug, which trains on numbers 255 times too large. Claim the range on the way in:

## The scores of images whose pixels must be in [0, 1].
fn spot(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 forward(img as Unit, w_e, tw1, tw2, cw1, cw2, w_h)

When claims fail

Each claim is checked where it’s made. Called from Python with the wrong data, every one of these raises an error that names the .ss file, line and column, and what was expected:

from_python(torch.rand(4, 1, 28, 28))  # not 16×16
blank = img.clone()
blank[2] = 0
inked(blank)                           # image 2 has no ink
spot(img * 255, *w)                    # bytes, not [0, 1]

ShapeError: spotter.ss:63:44: expected [b 1 S S] = (4, 1, 16, 16), got shape (4, 1, 28, 28)
ShapeError: spotter.ss:66:35: expected [b 1 S S] = (4, 1, 16, 16), got shape (3, 1, 16, 16)
DomainError: spotter.ss:71:20: expected Unit (values in [0, 1]), got values in [0, 249]

The second message is the claim of inked failing: nonblank kept 3 images, not b = 4. The third is a DomainError, a check of the values rather than the shape.

as versus to

as never changes a value; it only states what the value is. to, which Chapter 8 covers, converts. img as Unit checks that the pixels are in [0, 1]; img to Unit, for bytes, divides them by 255. Keep the two apart:

Turning checks off

The checks are cheap, but not free, so shapescript build and the Python import hook let you choose:

debug means “on, unless Python runs with -O”. So python -O would let spot(img * 255, *w) through, and the shape checks would still run. From Python, shapescript.install(checks="always", value_checks="off") sets them for imported files.

On a GPU, value checks don’t wait for the result: they use torch._assert_async, so a failure shows up at a later CUDA call as a device-side assert. To see which check failed, run the same inputs on the CPU.

Seeing the generated code

shapescript build writes the Python that a file compiles to. It’s worth reading once, to see that there’s no magic:

shapescript build spotter.ss -o build/spotter.py

The function inked becomes:

def inked(img):
    """The images, which must all have ink."""
    _ss_dev = _ss_rt.device()
    _ss_rt.check_rank(img, 4, "inked", "img", "[b 1 S S]", "spotter.ss:66:10")
    _ss_rt.check_dtype(img, _ss_torch.float32, "f32", "inked", "img", "[b 1 S S]", "spotter.ss:66:10")
    _ss_rt.check_device(img, _ss_dev, "inked", "img", "spotter.ss:66:10")
    _ss_d_b = img.shape[0]
    _ss_rt.check_sizes(img, (None, 1, S, S), "inked", "img", "[b 1 S S]", "spotter.ss:66:10")
    return _ss_check(nonblank(img), (_ss_d_b, 1, S, S), "spotter.ss:66:35", "[b 1 S S]")

The first four checks are on the way in: the number of axes, the kind of values (floats here, since the type gives no dtype and so means f32; bytes or booleans would be an error), the device (Chapter 16), and the sizes, with None for the generic b, whose size is read from the argument. The as is the _ss_check on the last line, which compares what nonblank returned with that b.

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

## 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 embed(img, w_e)
        |> mix_patches(tw1, tw2)
        |> mix_channels(cw1, cw2)
        |> norm
        |> head(w_h)

## Images of any size from Python, which must be S × S.
fn from_python(img: [b 1 h w]) -> Images = img as Images

## The images, which must all have ink.
fn inked(img: Images) -> Images = nonblank(img) as [b 1 S S]

## The scores of images whose pixels must be in [0, 1].
fn spot(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 forward(img as Unit, w_e, tw1, tw2, cw1, cw2, w_h)

Exercises

1. Write crop(img: [b 1 h w]) -> Images, which keeps the top-left S×S corner of each image, for photos larger than 16×16. A slice is img[:, :, :S, :S]; why does it need an as?

2. Make blend check without as, by changing only its signature. When is that the better fix?

   dim S = 16
   dim Images = [b 1 S S]
   
   fn blend(a: Images, c: [m 1 S S]) -> Images = (a + c) / 2
   

3. A colleague puts x = x as [b N D] after every stage of forward, “to be safe”. What does that cost, and what would you do instead?