Introduction

Here is a function in ShapeScript:

## Model width.
dim D = 768
## Number of attention heads.
dim H = 12
## Width of each head.
dim Dh = D / H

fn causal_attention(x: [b n D], w_qkv: [D 3*D], w_o: [D D]) -> [b n D]:
    q, k, v = x @ w_qkv
        |> split(b n (3 H Dh) -> b H n Dh)

    att = einsum(q, k: b H n Dh, b H m Dh -> b H n m)
        |> scale(1 / sqrt(Dh))
        |> masked_fill(!causal[n m], -inf)
        |> softmax(over=m)

    return einsum(att, v: b H n m, b H m Dh -> b H n Dh)
        |> rearrange(b H n Dh -> b n (H Dh))
        |> matmul(w_o)

If you’ve written attention in PyTorch, you can probably read it already. The differences are what make it checkable:

• Types have axes. x: [b n D] says x has three axes: a batch b and a length n, whose sizes the caller chooses, and D, which is 768.
• Patterns are syntax, not strings. split(b n (3 H Dh) -> b H n Dh) is einops’ notation, but the checker reads it: it knows 3*H*Dh must equal the last axis of x @ w_qkv, which is 3*D, and it knows that D / H divides evenly.
• Axes are names in scope. m is introduced by the first einsum, and from then on softmax(over=m) and causal[n m] mean that axis. There is no dim=-1.
• |> pipes a value into the next operation, so a computation reads top to bottom.

If you write softmax(over=k), pass a w_o of the wrong size, or write -> b H m n in the last einsum, the checker tells you where and why, before anything runs. When it’s satisfied, shapescript build turns the file into ordinary Python that calls torch and einops, with a runtime check of the arguments’ shapes on entry to every function.

Who this book is for

You should know some Python and some PyTorch. You don’t need to know einops, but if you do, Chapter 5 will feel familiar.

How to use this book

Read it in order: each chapter builds on the ones before, and most of the book is spent building programs.

• Chapters 1 and 2 set you up and write a first program, a small classifier.
• Chapters 3 to 11 are the language, taught by building one project: a shape spotter, a small image classifier that tells squares, disks, pluses and crosses apart. Each chapter adds what the spotter needs next, and ends with the whole of spotter.ss as it stands, so you can check yours against it. By Chapter 10 it’s trained, and gets 96% of new images right.
• Chapter 12 turns the spotter’s weights into a model and moves its training loop, with checkpoints and logs, into ShapeScript. Chapters 13 to 16 are the tools around it (the editor, notebooks, exporting and several devices), and Chapter 17 is a second project: a small transformer trained end to end.

Type the code in as you go. Every chapter ends with exercises, most of which extend the program you’re building: a fifth shape, another patch size, a new function. Try them before opening the solution. Each solution is folded away like this:

Many examples show what shapescript check prints. Errors are meant to be read, so the book shows them in full, exactly as the checker prints them (the tests compare them). Each diagnostic has a code, like E0401, and shapescript explain E0401 prints the section of the language reference that explains it.

Status

ShapeScript is in early development. What this book teaches works today: checking, compiling to Python, the editor support, notebooks, training, exporting, and training on several GPUs of one machine, which has run on three GPUs with the same results as one. Some things are still coming, such as training across several machines. The book marks each one where it comes up. ShapeScript itself is published, as the beta 0.1.4b1: pip install shapescript from PyPI, and the ShapeScript extension from the VS Code Marketplace. Chapter 1 takes you through installing them step by step, on macOS, Linux or Windows.

Source code

The book’s source is in the book/ folder of the ShapeScript repository. To read it locally, install mdBook and serve it:

cargo install mdbook --version 0.5.4 --locked
mdbook serve book --open