Value Domains and Dtypes

The spotter’s shapes are right, but the checker only knows half of what its tensors are. To it, every number so far is an f32 that could be anything. That isn’t true of the data the spotter really works with:

• An image file stores each pixel as a byte, a whole number from 0 to 255, not as an f32.
• A label is a class number: 0 for a square, 1 for a disk, 2 for a plus, 3 for a cross. It isn’t a measurement.
• head returns scores you can turn into probabilities, but they aren’t probabilities yet.

Mixing these up gives a second family of bugs, about values: probabilities passed where logits belong, a boolean mask added to numbers, log of something that can be zero, labels used to index the wrong table. PyTorch runs all of these. In this chapter the spotter learns to take bytes and labels, and gets a loss and a prediction. Along the way, ShapeScript’s domains catch each of these mistakes.

The element type

A tensor type is a shape, then an optional domain, then an optional dtype:

[b 1 S S] u8           # dtype only: bytes, as images are stored
[b C] Logit over C     # domain only: class scores over C
[b] Index[C]           # domain only: one class number per image
[b n D] Signed bf16    # both: values in [-1, 1], stored as bfloat16

The domain says which values are allowed, and the dtype says how they’re stored. With no domain, it’s Real, any finite number. With no dtype, it’s the one the domain implies: bool for masks, i64 for indices and f32 for the rest.

From bytes to numbers

Images arrive as bytes, so forward should take Images u8. Neural networks work best when their inputs are centered on zero, so the first thing forward does is turn the bytes into numbers in [-1, 1]. Two conversions do it, written with to: bytes to Unit, which divides by 255, then Unit to Signed, which computes 2x − 1. In spotter.ss, change forward, and give head’s scores their domain while you’re there:

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

## The class scores of each image.
fn forward(img: Images u8, w_e: [P*P D], tw1: [N Dt], tw2: [Dt N],
           cw1: [D Dc], cw2: [Dc D], w_h: [D C]) -> [b C] Logit over C:
    return img
        |> to Unit
        |> to Signed
        |> embed(w_e)
        |> mix_patches(tw1, tw2)
        |> mix_channels(cw1, cw2)
        |> norm
        |> head(w_h)

Images u8 is the alias with a dtype after it: [b 1 S S] u8. Forget the conversions and pass the bytes straight to embed, and it’s an error, because embed takes f32 numbers. A parameter without a dtype means f32, so a function says what it’s given, down to how it’s stored.

While you’re in spotter.ss, delete Chapter 7’s from_python, inked and spot. They were experiments with as, and spot’s claim that its images are in [0, 1] no longer fits a forward that takes bytes.

Domains are inferred

You rarely write a domain inside a function. The checker tracks the range of every value through arithmetic and activations, so most domains follow from the code. After to Signed, the pixels are in [-1, 1]; after gelu, they’re at least about −0.17; norm divides by a square root, which is never negative.

Inference catches mistakes too. Say you want each image’s ink level, its average brightness in [0, 1], to find faint drawings. You copy the first two stages of forward, and average:

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

## The average brightness of each image.
fn ink_level(img: Images u8) -> [b] Unit:
    return img |> to Unit |> to Signed |> reduce(b c h w -> b, mean)

$ shapescript check ink_level.ss
error[E0501]: expected `Unit`, got `Signed`
 --> ink_level.ss:6:12
  |
6 |     return img |> to Unit |> to Signed |> reduce(b c h w -> b, mean)
  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ this is `[b] Signed`, but the return type is `[b] Unit`
  |
  = help: convert it with `to Unit`

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

The average of numbers in [-1, 1] is in [-1, 1], so it isn’t Unit. The pipeline should stop after to Unit.

The ranges are approximate: to the checker, x - x on Signed values is [-2, 2], not 0. When it can’t prove something you know, use as (Chapter 7).

Scores, probabilities and labels

Logit over C, Prob over C and LogProb over C say what a tensor means along the axis C. head returns Logit over C: unnormalized scores, one per class. A plain Real result, such as pool(x) @ w, counts as a logit, because it is one, so head needed no conversion. softmax(z, over=C) turns logits into Prob over C.

Now the spotter can learn. Training needs a loss: a number that’s small when the scores favor the right class. Add it to spotter.ss, with a function that gives each image’s likeliest class:

## 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: [N Dt], tw2: [Dt N],
        cw1: [D Dc], cw2: [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: [N Dt], tw2: [Dt N],
           cw1: [D Dc], cw2: [Dc D], w_h: [D C]) -> [b] Index[C]:
    return argmax(forward(img, w_e, tw1, tw2, cw1, cw2, w_h), over=C)

y: [b] Index[C] is the labels, one class number per image. -> [] NonNeg is a scalar that’s never negative, which cross-entropy never is.

cross_entropy takes logits and applies log_softmax itself. A common mistake is to apply softmax first because the loss “wants probabilities”. The model still trains, but badly, and nothing in PyTorch complains. ShapeScript knows what softmax returns:

dim C = 4

fn loss_of(z: [b C] Logit over C, y: [b] Index[C]) -> [] NonNeg:
    p = softmax(z, over=C)
    return cross_entropy(p, y, over=C)

$ shapescript check softmax_loss.ss
error[E0501]: expected `Logit over C`, got `Prob over C`
 --> softmax_loss.ss:5:26
  |
5 |     return cross_entropy(p, y, over=C)
  |                          ^ this is `[b C] Prob over C`
  |
  = help: softmax applied before cross_entropy?

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

The same rule catches softmax applied twice. Probabilities aren’t logits, so a second softmax is an error too.

Masks

Comparisons give a Mask, a tensor of booleans. img > 0.5 says which pixels have ink. To count them, you might add the mask to something, or multiply by it, as you would in PyTorch. Masks aren’t numbers, though, so arithmetic on them is an error, with a suggestion:

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

## The image with its inked pixels made one brighter.
fn boost(img: Images) -> Images:
    return img + (img > 0.5)

$ shapescript check ink_mask.ss
error[E0505]: cannot add `Mask` to `Real`
 --> ink_mask.ss:6:12
  |
6 |     return img + (img > 0.5)
  |            ^^^^^^-----------
  |                  |
  |                  this is `[b 1 S S] Mask`
  |
  = help: did you mean `to Additive`?

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

Either convert the mask to the numbers you mean, or use it as a mask. to Binary turns it into floats that are 0 or 1, and their mean is the fraction of the image that’s ink:

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

## The fraction of each image's pixels that have ink.
fn ink_fraction(img: Images) -> [b] Unit:
    return img > 0.5 |> to Binary |> reduce(b c h w -> b, mean)

## How many of each image's pixels have ink.
fn ink_count(img: Images) -> [b] i64 = reduce(img > 0.5, b c h w -> b, sum)

The checker knows the mean of zeros and ones is in [0, 1], so ink_fraction is Unit. Summing a mask counts its true elements, and gives whole numbers, i64.

The suggestion in the error, to Additive, is for attention, which the spotter doesn’t have. causal[n m] is the mask of which positions a token may look at, and to Additive turns it into 0 where it’s true and −∞ where it’s false, ready to add to attention scores:

fn masked(att: [b H n m]) -> [b H n m]:
    return att + (causal[n m] |> to Additive)

fn filled(att: [b H n m]) -> [b H n m]:
    return att |> masked_fill(!causal[n m], -inf)

Chapter 17 uses this to build a transformer.

Indices

Index[V] is a whole number that’s a position along an axis of size V. It connects values to shapes. The spotter’s labels are Index[C], and so are its predictions, since argmax, argmin, argsort and topk give positions along the axis they work on.

Say you want to draw each prediction in its class’s color, from a table of four RGB colors. Indexing the table with the predictions looks up a row per image:

dim C = 4
dim K = 5

## Each image's predicted class, as a color.
fn colors(pred: [b] Index[C], palette: [C 3] u8) -> [b 3] u8 = palette[pred]

fn colors_bad(pred: [b] Index[C], palette: [K 3] u8) -> [b 3] u8 = palette[pred]

$ shapescript check colors.ss
error[E0507]: expected `Index[K]`, got `Index[C]`
 --> colors.ss:7:76
  |
7 | fn colors_bad(pred: [b] Index[C], palette: [K 3] u8) -> [b 3] u8 = palette[pred]
  |                                                                    ------- ^^^^ this is `[b] Index[C]`
  |                                                                    |
  |                                                                    the axis it indexes is `K` = 5
  |
  = note: the sizes differ: `C` = 4, not `K` = 5

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

The second function uses a palette with five rows. That might be left over from a version of the spotter with five shapes. It would still run in PyTorch, since four classes fit in five rows, and it would draw the wrong colors. The checker makes sure an index is used with the axis it’s a position on.

The same rule keeps token ids with their vocabulary: emb[ids], with emb: [V D] and ids: [b n] Index[V], looks up embeddings, and ids from another vocabulary are E0507.

Converting with to

as claims; to converts. You’ve used three conversions. These are all of them:

Any other to is an error that lists these: there’s no conversion from bytes straight to Signed, so forward takes two steps. Converting to a dtype only, as in x to f16 or mask to f32, changes the storage and keeps the values.

Boundary warnings

Some functions have edges: log of 0 is −∞, and dividing by 0 gives inf. The checker warns when a value may sit on an edge, and refuses one that’s surely past it. The log-probabilities of the spotter’s classes are a natural thing to want, and the obvious way to write them is on an edge:

dim C = 4

fn log_probs(z: [b C]) -> [b C] = log(softmax(z, over=C))

$ shapescript check log_probs.ss
warning[W0501]: `softmax(z, over=C)` can be exactly 0
 --> log_probs.ss:3:39
  |
3 | fn log_probs(z: [b C]) -> [b C] = log(softmax(z, over=C))
  |                                       ^^^^^^^^^^^^^^^^^^ `log(0)` is -inf
  |
  = note: it's `Prob over C`, so its values are in [0, 1]
  = help: use `log_softmax` instead of `softmax`

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

When one score is far above the others, the others’ probabilities underflow to exactly 0 in floating point, and their log is −∞. That’s the kind of bug that shows up as a NaN loss three hours into training. log_softmax(z, over=C) computes the same thing without the problem, and gives LogProb over C.

Your own domains

A domain declaration names a range. Dividing the scores by a temperature before the softmax makes the spotter less sure of itself (above 1) or more sure (below 1), which is a common way to make its probabilities honest. A temperature must be positive, and silly values aren’t useful, so give it a domain:

## Sampling temperature.
domain Temp = [0.05, 10]

dim C = 4

## The class probabilities, softened or sharpened by `t`.
fn calibrated(z: [b C] Logit over C, t: Temp) -> [b C] Prob over C = softmax(z / t, over=C)

t: Temp is a shapeless parameter, so it takes a tensor of any shape whose values are in the range, or a number, as in calibrated(z, 1.5). Because the range excludes 0, the division gets no warning. ( and ) exclude an end, as in (0, 1].

Dtypes

The dtypes are bool, u8, i8, i16, i32, i64, f16, bf16, f32 and f64. Promotion follows PyTorch’s: [b D] bf16 * 0.5 stays bf16, and bf16 + f16 is f32.

To save memory, you might store the patches as bf16 and keep the weights in f32. Matrix products need one dtype, since torch doesn’t promote them:

dim S = 16
dim P = 4
dim N = (S / P) * (S / P)
dim D = 32

fn embed_half(patches: [b N P*P] bf16, w_e: [P*P D]) -> [b N D] = patches @ w_e

$ shapescript check dtypes.ss
error[E0502]: `@` takes tensors of one dtype, but these are `bf16` and `f32`
 --> dtypes.ss:6:67
  |
6 | fn embed_half(patches: [b N P*P] bf16, w_e: [P*P D]) -> [b N D] = patches @ w_e
  |                                                                   -------^^^---
  |                                                                   |         |
  |                                                                   |         this is `f32`
  |                                                                   this is `bf16`
  |
  = note: torch doesn't promote the dtypes of a matrix product
  = help: convert one of them, as in `to bf16`

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

Convert one side with to: (patches to f32) @ w_e, or patches @ (w_e to bf16) with a bf16 result.

The spotter so far

forward now takes bytes, and the spotter has a loss and a prediction:

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] 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: [N Dt], tw2: [Dt N],
           cw1: [D Dc], cw2: [Dc D], w_h: [D C]) -> [b C] Logit over C:
    return img
        |> to Unit
        |> to Signed
        |> embed(w_e)
        |> mix_patches(tw1, tw2)
        |> mix_channels(cw1, cw2)
        |> 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: [N Dt], tw2: [Dt N],
        cw1: [D Dc], cw2: [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: [N Dt], tw2: [Dt N],
           cw1: [D Dc], cw2: [Dc D], w_h: [D C]) -> [b] Index[C]:
    return argmax(forward(img, w_e, tw1, tw2, cw1, cw2, w_h), over=C)

Exercises

1. Add probs to the spotter: the probability of each class for each image, [b C] Prob over C, from the same weights as predict.

2. Add confidence, which takes the scores z: [b C] and gives the probability of each image’s likeliest class. What domain can it promise?

3. A fifth shape joins the data: a triangle. What do you change in spotter.ss, and which types change with it? What happens to colors in this chapter, with its [C 3] palette, and to colors_bad?

4. This should be each image’s average brightness. Why is it an error, and what’s the fix, given that the images are already in [0, 1]?

   dim S = 16
   dim Images = [b 1 S S]
   
   fn brightness(img: Images) -> [b] Unit = reduce(img to Unit, b c h w -> b, mean)