Modules, Imports and Packages

spotter.ss has grown to about a hundred lines. Most of them are about the spotter: its sizes, its patches, its classes. But the mixer blocks, norm, mix_patches and mix_channels, aren’t. Any model that works on tokens could use them. This chapter moves them into a file of their own, and shows how files find each other, and how to share a file of ShapeScript with other people.

Each .ss file is a module with its own names, and import brings in another file. Packages of modules are shared the way Python’s are: with pip.

A file for the mixer blocks

Make a file mixer.ss next to spotter.ss, and move the three blocks into it. On the way, make them generic. In spotter.ss they took [b N D], the spotter’s own sizes; here they take [b n d], whatever the caller has (Chapter 4). A block that names no dims fits any model:

## Layer normalization over the last axis, 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)

## Mixes information between the tokens.
fn mix_tokens(x: [b n d], w1: [n k], w2: [k n]) -> [b n d]:
    h = einsum(norm(x), w1: b n d, n k -> b k d) |> gelu
    return x + einsum(h, w2: b k d, k n -> b n d)

## Mixes information between the channels of each token.
fn mix_channels(x: [b n d], w1: [d k], w2: [k d]) -> [b n d]:
    return x |> norm |> matmul(w1) |> gelu |> matmul(w2) |> add(x, _)

mix_patches is now mix_tokens, since outside the spotter the tokens needn’t be patches. Its weights’ hidden axis is k, not h, which needs a word of explanation; we’ll come back to it below.

Importing a file

spotter.ss imports the new file, and uses its names qualified with the file’s name:

devices cpu

import python torch
import mixer

and forward calls the blocks as mixer.mix_tokens, mixer.mix_channels and mixer.norm:

## 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 |> mixer.mix_tokens(tw1[l], tw2[l]) |> mixer.mix_channels(cw1[l], cw2[l])
    return x |> mixer.norm |> head(w_h)

Delete norm, mix_patches and mix_channels from spotter.ss, and the rest is unchanged. The calls are checked exactly as before: at mixer.mix_tokens(tw1[l], tw2[l]), n is N, d is D and k is Dt, and the result is [b N D] again, with the spotter’s names.

Everything a file declares can be used this way: functions, dims, aliases, and the models and records of Chapter 12. There’s no from mixer import norm: a name from another file is always written with the file’s name, so you can see where it comes from.

import blocks.mixer imports blocks/mixer.ss and binds mixer; import blocks.mixer as mx binds mx, so the calls become mx.norm(x).

Dims belong to their file

Why make the blocks generic, rather than move the dims to mixer.ss too? Because another file’s dims keep their file’s name. Say layers.ss declared a width of its own:

## Token width.
dim D = 32

fn norm(x: [b n D]) -> [b n D] = layer_norm(x, over=D)

layers.D is a different dim from the spotter’s D, even though both are 32, just as two dims of one file are. Sizes still match, so a call works. Names don’t, so broadcasting by name doesn’t (Chapter 6):

import layers

dim D = 32

fn shift(x: [b n D], bias: [b n layers.D]) -> [b n D] = x + bias

$ shapescript check two_widths.ss
error[E0404]: can't line up the axes of `[b n D]` and `[b n layers.D]`
 --> two_widths.ss:5:57
  |
5 | fn shift(x: [b n D], bias: [b n layers.D]) -> [b n D] = x + bias
  |                                                         -^^^----
  |                                                         |   |
  |                                                         |   this is `[b n layers.D]`
  |                                                         this is `[b n D]`
  |
  = note: neither operand's names are all in the other, so the axes line up from the right, and `D` meets `layers.D`
  = help: give them the same names with `rearrange` or `as`

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

If you mean them to be the same, an as says so: bias as [b n D]. That’s why mixer.ss has no dims: generic axes take the caller’s names, so its blocks fit into the spotter without an as.

Where files are found

An import looks next to the importing file first, then on the Python path, then in the standard library, then among installed packages. A file it can’t find is an error:

import mixr

$ shapescript check missing.ss
error[E2301]: there's no `mixr.ss` to import
 --> missing.ss:1:8
  |
1 | import mixr
  |        ^^^^
  |
  = note: an `import` looks next to the importing file, then on the Python path

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

Files can’t import each other in a circle (E2302): if mixer.ss imported spotter.ss, neither could be checked first. The help says what to do, and it’s what we just did: move what both need into a third file that both import.

Errors in an imported file

An imported file is checked with the importing one, and an error inside it is reported at the import that leads to it. Here is what happened when this chapter was first written, with the hidden axis of mix_tokens named h, like the value it multiplies:

fn mix_tokens(x: [b n d], w1: [n h], w2: [h n]) -> [b n d]:
    h = einsum(x, w1: b n d, n h -> b h d) |> gelu
    return x + einsum(h, w2: b h d, h n -> b n d)

import mixer_h

$ shapescript check uses_mixer_h.ss
error[E2303]: `mixer_h.ss` has 2 errors
 --> uses_mixer_h.ss:1:1
  |
1 | import mixer_h
  | ^^^^^^^^^^^^^^
  |
  = note: mixer_h.ss:2:5: error[E0303]: `h` is already an axis, so it can't also name a value
  = note: mixer_h.ss:3:30: error[E0403]: the pattern has 3 axes, but its input has 0

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

The notes give each error’s place in mixer_h.ss, and checking that file itself shows them in full. The first is the real one. Axes are names in scope (Chapter 5): w1: [n h] makes h the name of an axis for the whole function, so it can’t also be a variable. The second follows from it: einsum got the axis h as its first operand, which has no axes. Renaming the axis to k fixes both.

The standard library

The standard library ships with the compiler. Its first module is shapescript.std.fit, a training loop, which Chapter 12 uses:

import shapescript.std.fit as std

Packages

mixer.ss is the kind of file you’d want in other projects too. pip is ShapeScript’s package manager. A ShapeScript package is an ordinary Python package whose folders hold .ss files, so it’s built, published and installed with the usual tools:

shapescript new mixer_blocks             # a package to start from
python -m build mixer_blocks             # a wheel and an sdist in mixer_blocks/dist
twine upload mixer_blocks/dist/*         # publish it to PyPI
pip install mixer_blocks                 # or a wheel file, a folder, git+https://...

shapescript new makes a folder with a pyproject.toml, a README and src/mixer_blocks/layers.ss, an example module, which you’d replace with mixer.ss. The pyproject.toml names ShapeScript’s build backend:

[build-system]
requires = ["shapescript>=0.4"]
build-backend = "shapescript.build"

which checks every .ss file when the wheel is built, so a package with type errors can’t be built, and adds a typed Python module for each, so Python editors know the types.

Once installed, a .ss file imports a package’s module by its dotted name, import mixer_blocks.mixer, and uses mixer.norm, mixer.mix_tokens and the rest with their full types. Python imports it as any module: from mixer_blocks.mixer import norm. shapescript check and the editor find installed packages through your Python environment, the active virtualenv or the interpreter the VS Code Python extension uses.

A package records the compiler version that built it. A package built by a newer ShapeScript than yours is error E2901 at the import, rather than a confusing failure on syntax your compiler doesn’t know. The full guide, with editable installs and versions, is docs/packages.md.

What we built

Over Chapters 3 to 11, the spotter went from a few dims to a trained classifier in two files:

• Chapter 3 gave it its sizes and the type of its input, Images, and the checker made sure 16 pixels split into 4-pixel patches.
• Chapter 4 wrote its first functions, norm, pool and head, and showed how a call works out b and n.
• Chapter 5 cut images into patches with rearrange and mixed the patches with einsum, with no transpose to get wrong.
• Chapter 6 put the blocks together with pipes, and explained the broadcasting by name norm had relied on all along.
• Chapter 7 ran it from Python for the first time, and showed the checks that guard every call.
• Chapter 8 said what its numbers mean: bytes in, logits out, labels as Index[C], and a loss.
• Chapter 9 stacked its blocks with a loop, and measured it with iterator chains.
• Chapter 10 trained it from Python, to 96.4% on images it hadn’t seen, loaded its weights back in ShapeScript, and gave it typed extern fns for nudging and shrinking images.
• Chapter 11 split it into a reusable module and the spotter itself.

Two things are still awkward. Every function takes the same six weight tensors, w_e, tw1, tw2, cw1, cw2 and w_h, in the same order, and nothing but care keeps them in it. And the training loop is Python, so its optimizer, its batches and its state are invisible to the checker. Chapter 12 fixes both: a model declares the weights once, and the training loop moves into ShapeScript.

The spotter so far

devices cpu

import python torch
import mixer

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

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

## A token per patch.
fn embed(img: Images, w_e: [P*P D]) -> [b N D] = patchify(img) @ w_e

## 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 |> mixer.mix_tokens(tw1[l], tw2[l]) |> mixer.mix_channels(cw1[l], cw2[l])
    return x |> mixer.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. Each mixer block of the spotter is a mix_tokens followed by a mix_channels. Add block(x: [b n d], tw1: [n k], tw2: [k n], cw1: [d j], cw2: [j d]) -> [b n d] to mixer.ss, which does both, and use it in the spotter’s loop.
2. Move per_class_accuracy and confusion out of spotter.ss, into a file metrics.ss that imports spotter for C. Then say what happens if spotter.ss also imports metrics, to report its accuracy.
3. The spotter’s head averages the tokens and multiplies by a matrix, which suits any token model. Make it generic over the token width and the number of classes, and move it to mixer.ss. What does forward’s return type [b C] Logit over C need now, if anything?