Functions and Generic Sizes

Functions are where shapes meet: a signature says what a function takes and gives, and every call is checked against it. In this chapter the spotter gets its first three functions: norm, which normalizes each token; pool, which averages an image’s tokens into one; and head, which turns that average into class scores. Along the way, we’ll see how to write functions, and how a call works out the sizes of a function’s generic axes.

Declaring functions

At the end of Chapter 3, spotter.ss held the sizes, the Images alias and nonblank. Each patch of an image will become a token, a row of D numbers, so a batch of images becomes a [b N D]. Mixing tokens works better when each token is normalized first: shifted to mean 0 and scaled to variance 1 over its channels. That’s layer normalization, and here it is written out. Add it to spotter.ss:

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

A function has a name, parameters, an optional return type and a body. The body is an indented block, whose names (mu, var) are local to it. mean(x, over=D) is a [b n]: one mean per token. Subtracting it from the [b n D] needs no unsqueeze; Chapter 6 explains why.

A body that’s a single expression can be written after = instead. Here is pool, the average token of each image:

## The average token of each image.
fn pool(x: [b n D]) -> [b D] = mean(x, over=n)

Functions can call each other in any order: pool could come before norm in the file, or call a function declared at the bottom.

A ## comment directly above a declaration documents it. Editors show it on hover, along with the signature.

Parameters

norm’s eps = 1e-5 is a parameter with a default. A parameter has a type, a default, or both:

Host scalars are the plain numbers of your program: counters, rates, flags. They live in Python, not on the GPU. A default must be a plain number or a constant computed from numbers and dims, as in scale = 1 / D.

Calls pass arguments by position or by name, as in Python: norm(x), norm(x, eps=1e-6).

Return types

The return type is optional. Without one, a function’s type is that of what it returns. With one, every return is checked against it, and a block must return on every path. Suppose we let the caller choose between averaging the tokens and taking the first one, and forget the other case:

dim D = 32

fn summary(x: [b n D], average: bool) -> [b D]:
    if average:
        return mean(x, over=n)

$ shapescript check missing_return.ss
error[E0307]: `summary` doesn't always return a value
 --> missing_return.ss:3:4
  |
3 | fn summary(x: [b n D], average: bool) -> [b D]:
  |    ^^^^^^^                               ----- the signature says it returns this
  |
  = note: a loop may run no times, and an `if` without `else` may not run

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

Write return types for functions others call: they’re documentation the checker keeps true, and errors in the body are then reported in the body, not at every call.

How calls work out sizes

Now the head: average the tokens, then multiply by a weight matrix that gives one score per class:

## The class scores of each image.
fn head(x: [b n D], w: [D C]) -> [b C] = pool(x) @ w

head calls pool. When you call a function, its generic axes become unknowns, and the checker solves for them from the arguments. Each argument’s axes are matched with its parameter’s axes by position. At this call, pool’s b and n are matched with head’s own b and n, and D must be 32, which it is. The result type [b D] becomes head’s [b D], and @ w makes it [b C].

The spotter will call head with a [b N D], the tokens of a batch of images. Then head’s n is N, 16. Names don’t have to match: only the sizes, and the number of axes, must agree.

When they don’t, the error is at the call, and says which argument and which axis. Here the head’s weights were stored the other way round, as [C D], the way torch.nn.Linear stores them:

dim D = 32
dim C = 4

fn pool(x: [b n D]) -> [b D] = mean(x, over=n)
fn head(x: [b n D], w: [D C]) -> [b C] = pool(x) @ w

fn scores(tokens: [b 16 D], w_h: [C D]) -> [b C]:
    return head(tokens, w_h)

$ shapescript check bad_call.ss
error[E0401]: `head` can't take this as `w`
 --> bad_call.ss:8:25
  |
5 | fn head(x: [b n D], w: [D C]) -> [b C] = pool(x) @ w
  |                     -------- parameter declared here
...
8 |     return head(tokens, w_h)
  |                         ^^^ expected `[D C]`, found `[C D]`
  |
  = note: axis 1: `C` = 4, not `D` = 32
  = note: axis 2: `D` = 32, not `C` = 4

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

In PyTorch, x @ w_h with a (4, 32) weight would fail too, but only when the line runs, and without saying which argument was wrong. Here it’s found before anything runs.

A tensor with the wrong number of axes is a different error, E0403. The tokens of a single image, without a batch axis, are a [N D], which pool can’t take:

fn one_image(tokens: [N D]) -> [D]:
    return pool(tokens)

What the result is called

A result axis keeps a name when the name means something to the caller. A dim keeps its name, and a generic axis takes the name of the argument’s axis it was matched with, if that axis had one. So pool of a [b N D] is a [b D], and of a [4 N D] it’s a [4 D]. This is how over= keeps working after a call: the axes are still called what you called them.

Inside the function

Inside its body, a function knows nothing about its generic sizes except what the signature says. The sizes are the caller’s. You can still use them as numbers: n in an expression is the size of the axis n, a Python int at run time, read from the argument’s shape. Here is pool with the mean written by hand, a sum divided by the number of tokens:

dim D = 32

## The average token of each image.
fn pool(x: [b n D]) -> [b D] = sum(x, over=n) / n

$ shapescript check pool_by_hand.ss
warning[W0501]: `n` can be exactly 0
 --> pool_by_hand.ss:4:49
  |
4 | fn pool(x: [b n D]) -> [b D] = sum(x, over=n) / n
  |                                                 ^ dividing by 0 gives inf or NaN
  |
  = note: it's `NonNeg`, so its values are in [0, inf)
  = help: if it's never 0, say so with `as Pos`

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

A warning, not an error: the program compiles and runs. The checker tracks the range of every number (Chapter 8), and as far as it can tell here, a size is a whole number that might be 0, so dividing by it might give inf. You know better: an image always has patches. The help says how to tell the checker:

dim D = 32

fn pool(x: [b n D]) -> [b D] = sum(x, over=n) / (n as Pos)

as Pos is an assertion, which Chapter 7 is about: a claim that the value is positive, which the checker then relies on and the program checks when it runs. In the spotter we keep mean(x, over=n), which does the same thing.

Unknown sizes in return types

A function that decides a size itself writes ?k in its return type, as nonblank did in Chapter 3. Here’s another: the tokens of one image whose channels aren’t all zero.

dim D = 32

fn live_tokens(x: [n D]) -> [?k D]:
    return x[sum(abs(x), over=D) > 0]

The body may return any number of rows, and each call’s result has a fresh unknown size. Chapter 7 shows how to state what that size is with as.

Elementwise functions

The spotter’s mixer blocks will need an activation function between their matrix products. We’ll use the built-in gelu, but you can write your own. A function whose parameters have no shape, only a domain such as Real, is elementwise: it applies to a tensor of any shape, element by element:

dim D = 32

## Swish, with a slope.
fn swish(x: Real, beta = 1.5) -> Real = x * sigmoid(beta * x)

fn tokens(h: [b n D]) -> [b n D] = swish(h)

fn scores(z: [b 4]) -> [b 4] = z |> swish(beta=2.0)

That’s how you write your own activation function. In the editor, hovering swish draws its curve (Chapter 13). A function of one tensor of any fixed shape that only does elementwise maths, such as fn relu2(x: [b D]) -> [b D] = square(relu(x)), gets a curve too.

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

Exercises

1. Real layer normalization also has a learned scale and shift per channel. Add g: [D] and beta: [D] parameters to norm, so that it returns the normalized tokens times g, plus beta.
2. Write max_pool(x: [b n D]) -> [b D], which keeps the largest value of each channel over the tokens, and head2, a version of head that concatenates the mean and the max pools before the weights. What shape must its weights have? Hint: cat(a, c, over=D) puts two tensors side by side along D.
3. Fix one_image above so that it checks, without changing pool.
4. summary above forgot a case. Make it return the first token when average is false. (x[:, 0] is the first token of each image.)