Shapes and Axes

In PyTorch a tensor’s shape is a tuple of numbers you can only see at run time. In ShapeScript it’s part of the tensor’s type, written in brackets, and each entry is an axis that can have a name. This chapter is about what can go in the brackets.

It’s also where our project starts. From here to Chapter 11, each chapter adds to one program, and teaches what that program needs next.

The project: a shape spotter

We’ll build a shape spotter: a classifier that looks at a small grayscale picture and says which of four shapes is drawn in it. Here is one of each, with # for ink:

square             disk               plus               cross
#############...   ................   ................   ................
#...........#...   ................   ................   ................
#...........#...   ......#.........   ................   ................
#...........#...   ....#####.......   .......#........   ................
#...........#...   ...#######......   .......#........   ................
#...........#...   ...#######......   .......#........   ...#.....#......
#...........#...   ..#########.....   ....#######.....   ....#...#.......
#...........#...   ...#######......   .......#........   .....#.#........
#...........#...   ...#######......   .......#........   ......#.........
#...........#...   ....#####.......   .......#........   .....#.#........
#...........#...   ......#.........   ................   ....#...#.......
#...........#...   ................   ................   ...#.....#......
#############...   ................   ................   ................
................   ................   ................   ................
................   ................   ................   ................
................   ................   ................   ................

Each picture is 16×16 pixels. The shapes come in different sizes and places, and the background is noisy, so the spotter can’t just look at a few fixed pixels.

The network is a small MLP-Mixer. It cuts each image into 16 patches of 4×4 pixels, and turns each patch into a token, a vector of 32 numbers. Then come two mixer blocks. Each block first mixes information between the patches, so a patch can learn what’s around it, and then within each token. Finally, it averages the tokens and turns the average into four scores, one per shape.

Every one of those steps changes the shapes of the tensors, which makes it a good program to learn shapes on. Chapter 7 runs the spotter from Python for the first time, and Chapter 10 trains it. After that, Chapter 12 shows how to write the training itself in ShapeScript.

Tensor types

A tensor type lists its axes:

[b 1 S S]    # four axes: a batch b, one channel, and two axes of size S
[16 16]      # two axes with fixed sizes
[]           # no axes: a scalar

Each axis is one of these:

After the shape, a type may give a domain and a dtype, as in [b C] Prob over C or [b 1 S S] u8. Chapter 8 covers those; until then, every tensor holds ordinary f32 numbers.

Dims

Make a folder called spotter, and in it a file spotter.ss. Start with the sizes. A dim declares a constant size:

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
## Number of classes: square, disk, plus and cross.
dim C = 4

(devices cpu says where the program runs; Chapter 16 is about it.) A dim’s value can use other dims, in any order, and arithmetic. Here N is worked out from S and P: 4 patches across and 4 down make 16.

Division must be exact. S / P is fine because 4 divides 16. Make the patches 5 pixels wide and the checker refuses, because 16 pixels can’t be cut into whole 5-pixel patches:

dim S = 16
dim P = 5
dim N = (S / P) * (S / P)

$ shapescript check bad_dims.ss
error[E0305]: `S` (16) isn't divisible by `P` (5)
 --> bad_dims.ss:3:9
  |
1 | dim S = 16
  |         -- `S` is 16
2 | dim P = 5
  |         - `P` is 5
3 | dim N = (S / P) * (S / P)
  |         ^^^^^^^ 16 / 5 isn't a whole number
  |
  = note: dims must be whole numbers

error[E0305]: `S` (16) isn't divisible by `P` (5)
 --> bad_dims.ss:3:19
  |
1 | dim S = 16
  |         -- `S` is 16
2 | dim P = 5
  |         - `P` is 5
3 | dim N = (S / P) * (S / P)
  |                   ^^^^^^^ 16 / 5 isn't a whole number
  |
  = note: dims must be whole numbers

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

It’s reported twice, once for each S / P. That’s a bug caught before any image is cut: in PyTorch, a 5-pixel patch size would either fail deep inside a reshape or quietly drop the last pixel of each row.

Writing D in a type, as in [b N D], does two things: it says the axis has size 32, and it names the axis D. Names matter, because operations find axes by name (softmax(x, over=C)), and because broadcasting lines axes up by name.

Sizes are compared after simplification, with every dim replaced by its value. So N, (S / P) * (S / P), 4*P and 16 are all the same size, and the checker accepts one where another is expected.

Generic axes

How many images does the spotter look at at once? It shouldn’t matter: one image, a batch of 64 for training, 500 for a test. A name in a parameter’s type that isn’t declared anywhere is a generic axis, and the caller chooses its size:

dim S = 16

## The average of a batch of images.
fn average(img: [b 1 S S]) -> [1 S S] = mean(img, over=b)

One call can pass a [64 1 S S], the next a [1 1 S S]. Inside the function, nothing is assumed about b. You can’t treat it as 64, and two different generic axes are never equal, because a caller could make them different. Subtracting one batch from another is an error:

dim S = 16

fn difference(x: [b 1 S S], y: [m 1 S S]) -> [b 1 S S]:
    return x - y

$ shapescript check generic_bad.ss
error[E0404]: can't line up the axes of `[b 1 S S]` and `[m 1 S S]`
 --> generic_bad.ss:4:12
  |
4 |     return x - y
  |            -^^^-
  |            |   |
  |            |   this is `[m 1 S S]`
  |            this is `[b 1 S S]`
  |
  = note: neither operand's names are all in the other, so the axes line up from the right, and `b` meets `m`
  = help: give them the same names with `rearrange` or `as`

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

b and m might be the same size at a particular call, but the function can’t rely on it. Chapter 4 shows how calls choose the sizes, and Chapter 7 how to tell the checker something it can’t prove.

Literal sizes

A number is an exact size. The spotter’s images are grayscale, so they have one channel, written 1. Literal axes have no name, so [b 1 S S] has named axes b and S, and one unnamed axis of size 1. Taking the channel axis away by its “name” doesn’t work:

dim S = 16

fn squeeze(img: [b 1 S S]) -> [b S S]:
    return mean(img, over=1)

over= needs an axis name. Unnamed axes are reached by position, through a pattern (Chapter 5):

dim S = 16

fn squeeze(img: [b 1 S S]) -> [b S S]:
    return img |> rearrange(b 1 h w -> b h w)

The pattern describes the four axes in order, checks the second has size 1, and leaves it out of the result.

Groups and grouped sizes

Cutting an image into 4×4 patches only works if its sides are multiples of 4. Our images are 16×16, and the dims already prove that 4 divides 16. But suppose we wanted to accept images of any size that can be cut, say 16×16 today and 32×24 tomorrow. A type can say “a multiple of 4”.

A group, (gh 4), is one axis whose size is the product of its parts. A grouped size, 4gh, a number touching a name, is shorthand for (gh 4):

## Cuts each image into 4×4 patches, one row of pixels per patch.
fn patchify(img: [b 1 4gh 4gw]) -> [b (gh gw) 16]:
    return img |> rearrange(b c (gh p1) (gw p2) -> b (gh gw) (p1 p2 c))

img: [b 1 4gh 4gw] takes an image whose height and width are multiples of 4, and introduces gh and gw, the number of patches down and across, as sizes of their own, which the return type uses. Calling it with a 16×16 image works out gh = gw = 4, so the result is [b 16 16]: 16 patches of 16 pixels. Calling it with an image 18 pixels high is an error at the call:

fn bad_call(img: [8 1 18 16]) -> [8 (gh gw) 16]:
    return patchify(img)

$ shapescript check patches.ss
error[E0301]: unknown size `gh`
 --> patches.ss:5:38
  |
5 | fn bad_call(img: [8 1 18 16]) -> [8 (gh gw) 16]:
  |                                      ^^
  |
  = help: a return type can only use sizes from the parameters and dims; for a size the function chooses, write `?gh`

error[E0301]: unknown size `gw`
 --> patches.ss:5:41
  |
5 | fn bad_call(img: [8 1 18 16]) -> [8 (gh gw) 16]:
  |                                         ^^
  |
  = help: a return type can only use sizes from the parameters and dims; for a size the function chooses, write `?gw`

error[E0402]: `patchify` can't take this as `img`
 --> patches.ss:6:21
  |
2 | fn patchify(img: [b 1 4gh 4gw]) -> [b (gh gw) 16]:
  |             ------------------ parameter declared here
...
6 |     return patchify(img)
  |                     ^^^ expected `[b 1 (gh 4) (gw 4)]`, found `[8 1 18 16]`
  |
  = note: axis 3: `(gh 4)` can't be 18, because 18 isn't divisible by 4

For more information about E0301 and E0402, run `shapescript explain` with the code.

Three errors. The last is the one we were after: 18 isn’t a multiple of 4. The first two say that gh and gw belong to patchify, so bad_call can’t use them in its own return type. Chapter 4 explains why.

The checker also confirms the arithmetic inside patchify, whose rearrange Chapter 5 explains: p1 = p2 = 4 and c = 1, so (p1 p2 c) is 4 × 4 × 1 = 16, as the return type says. Return [b (gh gw) 32] instead, and the function itself is an error.

4 gh with a space is two axes, 4 and gh. When that’s probably a typo, the checker gives a hint:

fn patches(img: [b 1 4 gh 4gw]) -> [b (gh gw) 16]:
    return img |> rearrange(b c p1 gh (gw p2) -> b (gh gw) (p1 p2 c))

$ shapescript check hint.ss
hint[H2201]: `4 gh` is two axes: did you mean `4gh`?
 --> hint.ss:1:22
  |
1 | fn patches(img: [b 1 4 gh 4gw]) -> [b (gh gw) 16]:
  |                      ^^^^ an axis of size 4, then the axis `gh`
  |
  = note: `gh` is part of a group elsewhere in the signature, but an axis of its own only here
  = help: for one axis of size 4 × `gh`, write `4gh`, without the space

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

The spotter itself only takes 16×16 images, so it doesn’t need grouped sizes: its type will say S, and the dims say the rest. Chapter 5 writes its patchify.

Shape aliases

The spotter will pass batches of images around a lot. A shape you write often can get a name: dim followed by brackets declares a shape alias. Add one to spotter.ss:

## A batch of grayscale images.
dim Images = [b 1 S S]

An alias can be a type’s whole shape, img: Images, or part of it: [Images D] would be [b 1 S S D]. The names inside it belong to the function that uses it, so each function that takes Images has its own generic axis b:

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

fn average(img: Images) -> [1 S S] = mean(img, over=b)

Unknown sizes

Some sizes can’t be known until the program runs. Some of the spotter’s input might be blank, with no shape drawn at all, and we might want to keep only the images with some ink: a pixel brighter than 0.5. How many there are depends on the data. Add this to spotter.ss:

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

reduce(img, b c h w -> b, max) is each image’s brightest pixel (Chapter 5 explains patterns), comparing it with 0.5 gives one true or false per image, and indexing with that keeps the images where it’s true.

?k is an unknown size. It’s a definite number, just one the checker can’t see, so it equals only itself. Every call to nonblank returns a new unknown size, so two calls give two different ?ks:

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

fn nonblank(img: Images) -> [?k 1 S S]:
    return img[reduce(img, b c h w -> b, max) > 0.5]

fn compare(x: Images, y: Images) -> [?k 1 S S]:
    return nonblank(x) - nonblank(y)

$ shapescript check unknown.ss
error[E0401]: mismatched shapes `[?k 1 S S]` and `[?k2 1 S S]`
 --> unknown.ss:8:12
  |
8 |     return nonblank(x) - nonblank(y)
  |            -----------^^^-----------
  |            |             |
  |            |             this is `[?k2 1 S S]`
  |            this is `[?k 1 S S]`
  |
  = note: `?k` is not known to equal `?k2`

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

Even with x and y the same batch size, the two filters can keep different numbers of images. ?k can’t appear in a parameter’s type: there, the caller chooses the size, so it’s a generic axis and is written k. Chapter 7 shows how as turns an unknown size into a known one, with a check at run time.

The spotter so far

Check it:

shapescript check spotter.ss

No output: it’s correct. There’s no model in it yet, only sizes and a type, but they already say a lot: the images are 16×16, they cut into 16 patches, and every token is 32 wide.

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

Exercises

1. The spotter’s patches could be smaller. Change P so that each image is cut into 64 patches. Which other dims change, and which values of P does the checker accept at all?
2. A colleague wants to spot shapes in 28×28 pictures, the size of MNIST digits, with the same 4×4 patches. Change the dims. How many patches does an image have now?
3. Declare an alias ColorImages for a batch of RGB images of the spotter’s size, and write fn to_gray(img: ColorImages) -> Images, which averages the three channels. You need a pattern, as in Literal sizes.
4. Write unpatchify, the inverse of the patchify with grouped sizes: it takes [b gh gw (4 4 1)], the patches with their counts down and across as two axes, and returns [b 1 4gh 4gw]. You need rearrange; mind the order of the parts of each group.