Foreword

Every PyTorch programmer knows this bug. A tensor that should be (batch, heads, seq, seq) is (batch, seq, heads, seq), a transpose went the wrong way, or a softmax(dim=-1) normalizes the wrong axis. Nothing complains. The model trains, a little worse than it should, and you find out days later, if at all.

The shapes were in your head the whole time. They were in the comments too: # x: (B, T, C). But nothing read the comments.

ShapeScript is a small language that reads them. Axis names and sizes are part of its types, every operation takes axes by name, and a checker verifies the whole program before any of it runs, in your editor, as you type. When it’s happy, it compiles to plain Python that calls PyTorch, so everything you know and every library you use still works.

This book teaches the language from the beginning. It assumes you’ve written some PyTorch: you know what a tensor is, what x @ w does and roughly what a training loop looks like. It doesn’t assume you’ve used einops, a typed language or a compiler.

Like The Rust Programming Language, which it borrows its shape from, the book is meant to be read in order, with your editor open. Every ShapeScript example in it is checked by the project’s test suite, so if the book says a program is fine, the checker agrees, and if the book shows an error, that is the error you’ll see.

ShapeScript is in early development. The book only teaches what works today. Where something is planned but not built yet, it says so, and calls it coming.

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

Getting Started

This chapter takes you from nothing to a working setup: ShapeScript installed, a project with its own Python environment, VS Code showing shapes as you type, and a notebook running ShapeScript cells. Then it checks and runs a first program and calls it from Python. Later chapters explain everything you type here.

What you’ll set up

ShapeScript has two parts:

• a compiler and a language server, written in Rust: the shapescript and shapescript-lsp programs;
• a Python package, shapescript, with the runtime, the import hook and the notebook magic.

Both come prebuilt. pip install shapescript installs the Python package with the two programs inside it, for Linux, macOS and Windows, and the VS Code extension, from the Marketplace, brings its own copy of the language server. You need neither Rust nor Node.js; Building from source, at the end of the setup, is for other platforms, for versions that aren’t released, and for working on ShapeScript itself.

The commands are given for the usual shells: bash and zsh on macOS and Linux, fish, and PowerShell on Windows. A block of commands is marked with its shell when that matters: bash, zsh, fish or PowerShell. An unmarked block works in any of them, or, when a PowerShell block follows it, is the macOS and Linux version. If something goes wrong, Appendix D has the usual problems and their fixes.

Prerequisites

Python 3.9 or later

The runtime and PyTorch are Python packages, so you need Python. Any version from 3.9 works; 3.12 or 3.13 are the safest choice, since PyTorch has wheels for them everywhere. Check what you have:

python3 --version

py --version

If it’s missing or older than 3.9, install one:

• macOS: with Homebrew, brew install python@3.13, which is /opt/homebrew/bin/python3.13; or the installer from python.org.
• Linux: your distribution’s package; on Debian and Ubuntu you also need its venv module: sudo apt install python3 python3-venv python3-pip. On Fedora, sudo dnf install python3.
• Windows: winget install Python.Python.3.13, or the installer from python.org. Use py to run it: on Windows, python and python3 can open the Microsoft Store instead of running Python.

On a Mac with Apple Silicon (M1 and later), make sure the Python you use is an Apple Silicon one:

python3 -c "import platform; print(platform.machine())"

It should print arm64. An x86_64 Python, such as one from an old Intel Homebrew in /usr/local/bin, runs under Rosetta, and recent PyTorch doesn’t install in it. If python3 is such a Python, name Homebrew’s by its full path, /opt/homebrew/bin/python3.13, in the commands below.

VS Code, with the Python and Jupyter extensions

Install VS Code (brew install --cask visual-studio-code on macOS, winget install Microsoft.VisualStudioCode on Windows), then Microsoft’s Python extension, which brings Pylance and the Python debugger, and Jupyter, for notebooks. From the Extensions view, or:

code --install-extension ms-python.python
code --install-extension ms-toolsai.jupyter

On macOS, the code command exists once you’ve run Shell Command: Install ‘code’ command in PATH from VS Code’s Command Palette (Cmd+Shift+P). On Windows and Linux, VS Code’s installer adds it.

A project with a virtual environment

Each project gets a virtual environment: a folder, .venv, with its own Python and its own packages. It keeps one project’s PyTorch from breaking another’s, and it’s how VS Code, notebooks and the command line agree on which packages you have: they all use the same .venv.

Create and activate it

Make a folder for the project, and the environment in it:

mkdir ~/hello-shapescript
cd ~/hello-shapescript
python3 -m venv .venv

mkdir ~\hello-shapescript
cd ~\hello-shapescript
py -3.13 -m venv .venv

(On a Mac where python3 isn’t Apple Silicon, /opt/homebrew/bin/python3.13 -m venv .venv.)

Then activate it. Activating puts .venv’s python, pip and, later, shapescript first on your PATH, in this terminal only, so do it in each new terminal you open for the project:

source .venv/bin/activate

source .venv/bin/activate.fish

.venv\Scripts\Activate.ps1

(In Windows’ older cmd.exe, it’s .venv\Scripts\activate.bat. If PowerShell says running scripts is disabled, see Appendix D.) Your prompt now starts with (.venv). deactivate turns it off. Check that python is the environment’s:

python -c "import sys; print(sys.prefix)"

It prints the path of your .venv. Then update pip, since an old pip can miss the wheels newer packages are published as:

python -m pip install --upgrade pip

Install PyTorch

Install PyTorch before ShapeScript. ShapeScript depends on it, and pip keeps the build you installed; if you install ShapeScript first, pip picks the default build for you.

PyTorch comes in builds for CPUs, NVIDIA GPUs (CUDA) and Apple GPUs. pytorch.org’s selector gives the exact command for your machine; the usual ones are:

• macOS, which includes Apple GPU support (mps), and Linux with an NVIDIA GPU, which gets CUDA (a download of a few GB):

  pip install torch numpy einops
  

• Linux or Windows without a GPU, a much smaller download:

  pip install torch --index-url https://download.pytorch.org/whl/cpu
  pip install numpy einops
  

• Windows with an NVIDIA GPU: the default build on Windows is CPU-only, so take the CUDA command from the selector, then pip install numpy einops.

ShapeScript needs einops. PyTorch warns Failed to initialize NumPy without numpy, and notebooks use it, so install it too. Check what you got:

python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

True means PyTorch sees a CUDA GPU. On a Mac, torch.backends.mps.is_available() tells you about the Apple GPU.

Install ShapeScript

With the environment active, install ShapeScript from PyPI:

pip install shapescript

The current release is the beta 0.1.4b1. pip downloads the wheel for your platform, with the compiler and the language server inside it, and einops with it; it keeps the PyTorch you installed. Check all three:

python -c "import shapescript; print(shapescript.__version__)"
shapescript --version
shapescript-lsp --version

The first prints 0.1.4b1, and the two programs 0.1.4-beta.1, Rust’s spelling of the same version. If pip says it found no version, your platform has no wheel (see Building from source), or your Python is older than 3.9.

Updating, and more projects

To update ShapeScript, run pip install -U shapescript in each project’s environment. VS Code updates the extension by itself.

For another project, make a new folder with its own .venv, and install PyTorch and ShapeScript into it as above.

To start a ShapeScript package, something others pip install, shapescript new mypkg makes one; Chapter 11 explains packages. To build one, pip fetches ShapeScript from PyPI into a separate build environment, and PyTorch with it, which is a large download. To build with the environment’s own ShapeScript instead:

pip install hatchling
pip install --no-build-isolation ./mypkg

Setting up VS Code

The VS Code extension shows errors as you type, the shape of anything you hover over, drawn as a diagram, and a lot more (Chapter 13). Setting it up takes four steps.

1. Install the extension

In VS Code’s Extensions view (Cmd+Shift+X on macOS, Ctrl+Shift+X elsewhere), search for ShapeScript, by the publisher shapescript, and click Install. Or from a terminal:

code --install-extension shapescript.shapescript --pre-release

The extension is published as a pre-release for now, which --pre-release asks for. Its language server comes inside it, built for your platform: Linux (x64, arm64), macOS (Apple Silicon, Intel) or Windows (x64, arm64).

The extension isn’t on Open VSX, the registry that VSCodium, Cursor and some other editors use. There, download the .vsix for your platform from the Marketplace page and install it with Install from VSIX… in the Extensions view’s ⋯ menu.

2. Open the project folder

Open the project’s folder, not just a file: File › Open Folder…, or from a terminal:

code ~/hello-shapescript

VS Code keeps its settings and its choice of Python per folder. If it asks whether you trust the folder’s authors, say yes: in Restricted Mode, extensions are limited.

3. Select the environment’s Python

Open the Command Palette (Cmd+Shift+P on macOS, Ctrl+Shift+P elsewhere), run Python: Select Interpreter, and pick the one in .venv (VS Code labels it ('.venv': venv), and often offers it by itself). Everything else uses this choice: Pylance looks up imports in its packages, the ShapeScript language server finds installed ShapeScript packages and Python sources through it, the debugger and notebooks run in it, and new terminals in VS Code activate it.

To make the choice part of the project, put it in .vscode/settings.json in the project folder:

{
  "python.defaultInterpreterPath": "${workspaceFolder}/.venv/bin/python"
}

On Windows, the path is ${workspaceFolder}\\.venv\\Scripts\\python.exe.

4. Check that it works

The extension finds its language server inside itself, so there’s nothing to configure. (Without a bundled server, as in an extension built another way, it uses the shapescript.server.path setting, which takes a full path, such as the environment’s .venv/bin/shapescript-lsp, and then shapescript-lsp on the PATH.)

Then, once you’ve typed hello.ss from Hello, ShapeScript below, check that:

• The status bar, bottom right, says ShapeScript, and keywords such as fn and dim are colored.
• Hovering x in normalize shows [n D] and a rectangle with its sides labeled n and D; hovering normalize shows its doc comment.
• Changing over=D to over=n underlines the call in red within a second, and the Problems panel lists E0501. Changing it back clears it.
• Once hello.ss is saved without errors, the Explorer shows an arrow beside it, and under it hello.py.
• In a Python file, from hello import normalize isn’t underlined, and hovering normalize shows its ShapeScript signature.

That hello.py is written by the extension, for Python’s sake. Pylance doesn’t know about .ss files, so in a Python file, from hello import normalize would be Import "hello" could not be resolved. The generated module declares normalize with its ShapeScript signature, so Pylance resolves the import, completes the name and shows the shapes on hover; and when Python imports it, it loads hello.ss. The Explorer nests it under the .ss file to keep it out of the way. A .py file you wrote yourself is never replaced.

If the status bar says Plain Text, or hovers show nothing, see Appendix D.

Debugging

To stop at breakpoints in .ss files, click left of a line number in fn main() to set one, then run ShapeScript: Debug File from the Command Palette, or from the ▷ menu at the top right of the editor. It uses the Python debugger that came with the Python extension, in the interpreter you selected. Without the Python extension, install the debugger into the environment, pip install debugpy, and name the environment’s Python in the setting shapescript.debug.python. Chapter 13 has more.

Other editors

shapescript-lsp is an ordinary language server, so any editor with LSP support can use it; with the environment active, it’s on your PATH, since the shapescript package installed it there. editors/neovim/shapescript.lua in the repository is a minimal Neovim setup.

Building from source

The published wheels and extension cover Linux, macOS and Windows on the usual processors. To run ShapeScript elsewhere, to try a commit that isn’t released, or to work on ShapeScript itself, build it from the repository instead. Everything else in this chapter stays the same: only Install ShapeScript and Install the extension change.

Prerequisites: Rust

You need Rust 1.88 or later. Install it with rustup, which also keeps it up to date.

On macOS and Linux:

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh

Rust needs a C linker. On macOS, install Apple’s command line tools with xcode-select --install; on Debian and Ubuntu, sudo apt install build-essential; on Fedora, sudo dnf install gcc.

On Windows:

winget install Rustlang.Rustup

On Windows, Rust links with Microsoft’s C++ build tools. If you don’t have Visual Studio, the installer offers to install them; or install them yourself with winget install Microsoft.VisualStudio.2022.BuildTools --override "--wait --passive --add Microsoft.VisualStudio.Workload.VCTools --includeRecommended".

Open a new terminal, so it sees the new PATH, and check:

cargo --version

In fish, if cargo isn’t found, add Rust’s folder to your PATH once (fish remembers it):

fish_add_path ~/.cargo/bin

If the version is older than 1.88, rustup update updates it.

Prerequisites: Git

You clone ShapeScript with Git. macOS has it once the command line tools are installed; on Linux, it’s the git package; on Windows, winget install Git.Git. Git for Windows also installs Git Bash, which you’ll use to build the VS Code extension.

Prerequisites: Node.js

Building the extension needs Node.js 20 or later: brew install node on macOS, winget install OpenJS.NodeJS.LTS on Windows, and nodejs.org’s packages on Linux (a distribution’s own can be older than 20). Check with node --version.

Clone and build

Clone the repository into your home folder and build the compiler and the language server. The commands are the same in every shell:

cd ~
git clone https://gitlab.com/NatanFreeman/shapescript.git
cd shapescript
cargo build --release -p shapescript-cli -p shapescript-lsp

The first build takes a few minutes. It writes shapescript and shapescript-lsp (with .exe on Windows) to target/release.

Now copy the two programs into the Python package, which is what the published wheels do. Installing the package then gives each virtual environment its own shapescript and shapescript-lsp commands, so you don’t need to put target/release on your PATH:

mkdir -p crates/py/python/shapescript/bin
cp target/release/shapescript target/release/shapescript-lsp crates/py/python/shapescript/bin/

New-Item -ItemType Directory -Force crates\py\python\shapescript\bin
Copy-Item target\release\shapescript.exe, target\release\shapescript-lsp.exe crates\py\python\shapescript\bin\

(Git ignores that folder.) If you also want the commands outside any virtual environment, for example for an editor other than VS Code, install them into Rust’s own folder, which rustup put on your PATH:

cargo install --path crates/cli --locked
cargo install --path crates/lsp --locked

shapescript run and shapescript export need a Python with PyTorch and ShapeScript. Run from a virtual environment’s shapescript, they use its Python; run from elsewhere, they use --python, or the SHAPESCRIPT_PYTHON environment variable, or python3 (Appendix D shows how to set it).

Install the package

In your project’s environment, after PyTorch, install the package from your clone instead of from PyPI:

pip install ~/shapescript/crates/py

pip install $HOME\shapescript\crates\py

pip compiles the package’s Rust part, which takes a minute, and installs it with the two programs you copied in. If shapescript --version then says this installation of shapescript has no shapescript binary, the copy step above was skipped: do it, and run pip install again.

If you’re working on ShapeScript itself, an editable install is handier: your changes to its Python code take effect without reinstalling (changes to its Rust need the command again). With the environment active:

pip install maturin
cd ~/shapescript/crates/py
maturin develop --release

To use ShapeScript, pip install is simpler: it doesn’t need maturin, and works inside conda, which maturin develop refuses (see Appendix D).

To update, run git pull in ~/shapescript, then the build, the copy and the pip install again, and rebuild the extension (below), so that the compiler, the runtime and the editor come from the same version.

Build the extension

The extension you build is a .vsix file, with the language server you built inside it. Name your platform: darwin-arm64 (a Mac with Apple Silicon), darwin-x64 (an Intel Mac), linux-x64, linux-arm64, win32-x64 or win32-arm64. On macOS and Linux:

cd ~/shapescript/editors/vscode
npm ci
cd ~/shapescript
ci/vsix.sh darwin-arm64
code --install-extension dist/vsix/shapescript-darwin-arm64.vsix

npm ci downloads the extension’s dependencies, and ci/vsix.sh builds it and packs it with target/release/shapescript-lsp. On Windows, run the same commands in Git Bash, with win32-x64 (or win32-arm64) for the platform: ci/vsix.sh is a bash script. You can also install the file from VS Code: in the Extensions view, the ⋯ menu has Install from VSIX…. Reload any open VS Code window afterwards (Developer: Reload Window in the Command Palette).

If you’re working on ShapeScript itself, you don’t need the .vsix: open the ~/shapescript folder in VS Code and press F5. VS Code builds the language server and the extension, then opens examples/ in a new window with the extension loaded.

The shapescript command

The command has a subcommand for each job. You’ll meet them all in this book:

shapescript check model.ss               # check types and shapes; print errors and warnings
shapescript explain E0401                # explain an error code
shapescript build model.ss -o model.py   # compile to Python
shapescript run program.py               # run a program that imports .ss files
shapescript export model.ss --fn f       # a torch.export program (Chapter 15)
shapescript export model.ss --hf out     # a Hugging Face model, for transformers and vLLM (Chapter 15)
shapescript new my_layers                # a package for pip (Chapter 11)

Hello, ShapeScript

In your project folder, create a file called hello.ss:

devices cpu

dim D = 4

## Turns each row of scores into probabilities.
fn normalize(x: [n D]) -> [n D] Prob over D = softmax(x, over=D)

fn main():
    x = 3 |> map(i => D |> map(j => j to f32))
    print("hello from ShapeScript")
    print(normalize(x))

Line by line:

• devices cpu says where the program’s tensors live. Every program says so once, and there’s no default. You’d write devices cuda[0] for a GPU, or devices mps on a Mac.
• dim D = 4 declares a constant size, D, which types can use.
• ## ... is a doc comment. Your editor shows it when you hover over normalize.
• fn normalize(x: [n D]) -> [n D] Prob over D declares a function that takes a tensor with two axes, any number n of rows and D columns, and returns one of the same shape whose values are probabilities that sum to 1 over D. The body after = is a single expression.
• fn main() is where a program run with shapescript run starts. Its body is indented, as in Python.
• 3 |> map(i => D |> map(j => j to f32)) builds a [3 D] tensor: for each of 3 rows, the indices 0 to 3 along D as floats. |> passes a value to the next step, and map computes its function, j => j to f32, for every index at once; Chapter 9 covers it.

Check it:

shapescript check hello.ss

It prints nothing: no news is good news. Now make a mistake. Change over=D to over=n, which is the wrong axis for normalizing rows:

devices cpu

dim D = 4

fn normalize(x: [n D]) -> [n D] Prob over D = softmax(x, over=n)

$ shapescript check hello_bad.ss
error[E0501]: expected `Prob over D`, got `Prob over n`
 --> hello_bad.ss:5:47
  |
5 | fn normalize(x: [n D]) -> [n D] Prob over D = softmax(x, over=n)
  |                                               ^^^^^^^^^^^^^^^^^^ this is `[n D] Prob over n`, but the return type is `[n D] Prob over D`
  |
  = help: convert logits with `to Prob over D`, or if these are already probabilities, say so with `as Prob over D`

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

Read the message from the top. It names the problem, points at the line, shows what was expected and what was found, and ends with a code you can look up. You’ll see many messages like it in this book.

Change it back and run the program:

shapescript run hello.ss

hello from ShapeScript
tensor([[0.0321, 0.0871, 0.2369, 0.6439],
        [0.0321, 0.0871, 0.2369, 0.6439],
        [0.0321, 0.0871, 0.2369, 0.6439]])

shapescript run runs a .ss file’s fn main(), or a Python program, in the Python it finds: the active virtualenv’s, or the one given with --python.

Calling ShapeScript from Python

Most of the time you’ll call ShapeScript from Python, from a training script or a notebook. A .ss file is a module Python can import, once the import hook is installed:

import torch
import shapescript
shapescript.install()          # lets `import` find .ss files

from hello import normalize    # compiles hello.ss, and caches the result
print(normalize(torch.randn(5, 4)))

The imported normalize is a Python function. It checks its argument when it’s called, so normalize(torch.randn(5, 3)) raises a shapescript.ShapeError naming hello.ss and the line, instead of computing garbage. Tracebacks point into the .ss file.

shapescript run program.py runs a Python program with the hook already installed, so a script run that way doesn’t need shapescript.install().

The editor also writes a small Python module beside each .ss file, hello.py beside hello.ss, with typed declarations of its functions, so Pylance and mypy know about them too. Chapter 13 has the details.

Notebooks

ShapeScript runs in Jupyter notebooks through a magic command. A notebook’s kernel is the Python that runs its cells, and it has to be your environment’s Python, which has ShapeScript. Install the kernel package into the environment, with it active:

pip install ipykernel

In VS Code, create a notebook (Create: New Jupyter Notebook in the Command Palette, or a new file ending in .ipynb), click Select Kernel at the top right, choose Python Environments, and pick .venv. (If you skipped pip install ipykernel, VS Code offers to install it.)

In Jupyter in the browser, install it into the same environment and start it from the project folder:

pip install jupyterlab
jupyter lab

Its kernel, Python 3 (ipykernel), is then the environment’s Python.

In the notebook’s first cell, load the extension once:

%load_ext shapescript

Then any cell that starts with %%shapescript is ShapeScript:

%%shapescript
devices cpu
dim D = 4
fn normalize(x: [n D]) -> [n D] Prob over D = softmax(x, over=D)

Run both. The second prints nothing: it only declares things, and they check. Chapter 14 is about notebooks.

Exercises

1. Run shapescript explain E0501. What does the reference say the code means?
2. In hello.ss, change the return type of normalize to -> [D n] Prob over D. Predict whether the checker accepts it, then check.
3. Change dim D = 4 to dim D = 0. What does the checker say?

A First Program: a Classifier

Let’s write something real: a small neural network that sorts points in the plane into three classes. We’ll write the model’s maths in ShapeScript and drive it from a short Python script, so you can see how the two fit together. Many of the ideas here get a chapter of their own later; this one is a tour.

Setting up

Make a folder called classifier, and in it a file classifier.ss. Start with the sizes:

devices cpu

## Number of input features.
dim D = 2
## Width of the hidden layer.
dim Hd = 32
## Number of classes.
dim C = 3

A dim is a size with a name. Types use these names, so the program says which 32 it means, not just 32.

The forward pass

The network takes a batch of points, multiplies by a weight matrix, applies a ReLU, and multiplies by another:

## The class scores of each example.
fn forward(x: [b D], w1: [D Hd], w2: [Hd C]) -> [b C] Logit over C:
    return x @ w1 |> relu |> matmul(w2)

Read the signature first:

• x: [b D] is a batch of points. b is a name that isn’t declared anywhere, so it’s a generic axis: the caller decides its size, and every call can use a different one.
• w1: [D Hd] and w2: [Hd C] are the weights, with fixed sizes.
• -> [b C] Logit over C is the result: one row of scores per point, one score per class. Logit over C says what the numbers mean: unnormalized scores over the axis C, the kind you’d pass to softmax or cross_entropy.

Then the body. x @ w1 is a matrix product, as in PyTorch: [b D] @ [D Hd] is [b Hd]. |> relu passes that to relu, and |> matmul(w2) passes the result as the first argument of matmul, so the line means matmul(relu(x @ w1), w2), which reads in the order it happens.

The checker follows every shape: [b D] @ [D Hd] is [b Hd], relu keeps it, and [b Hd] @ [Hd C] is [b C], which is what the signature promised.

The loss and predictions

## The mean cross-entropy of the scores against the labels.
fn loss(x: [b D], y: [b] Index[C], w1: [D Hd], w2: [Hd C]) -> [] NonNeg:
    return cross_entropy(forward(x, w1, w2), y, over=C)

## The likeliest class of each example.
fn predict(x: [b D], w1: [D Hd], w2: [Hd C]) -> [b] Index[C]:
    return argmax(forward(x, w1, w2), over=C)

y: [b] Index[C] is one label per point, and Index[C] says each is a whole number from 0 to C - 1, a position along the axis C. -> [] NonNeg is a scalar, a tensor with no axes, that is never negative.

cross_entropy(..., over=C) and argmax(..., over=C) take the class axis by name. There’s no dim=1 to get wrong: if the axis isn’t called C, it’s an error.

Check the file:

shapescript check classifier.ss

No output: it’s correct.

Training it from Python

Create train.py next to it:

import torch
from classifier import loss, predict

# Three blobs of points in the plane, one per class.
g = torch.Generator().manual_seed(0)
centers = torch.tensor([[0.0, 2.0], [-2.0, -1.0], [2.0, -1.0]])
y = torch.randint(0, 3, (300,), generator=g)
x = centers[y] + 0.7 * torch.randn(300, 2, generator=g)

w1 = (0.5 * torch.randn(2, 32, generator=g)).requires_grad_()
w2 = (0.5 * torch.randn(32, 3, generator=g)).requires_grad_()

for step in range(201):
    l = loss(x, y, w1, w2)
    l.backward()
    with torch.no_grad():
        for w in (w1, w2):
            w -= 0.1 * w.grad
            w.grad = None
    if step % 50 == 0:
        print(f"step {step}: loss {l.item():.3f}")

accuracy = (predict(x, w1, w2) == y).float().mean().item()
print(f"accuracy: {accuracy:.1%}")

and run it:

$ shapescript run train.py
step 0: loss 2.059
step 50: loss 0.094
step 100: loss 0.081
step 150: loss 0.077
step 200: loss 0.074
accuracy: 97.3%

from classifier import loss, predict compiled classifier.ss to Python, and the functions it imported are ordinary PyTorch code, so autograd, .backward() and everything else work as usual. (Chapter 12 writes this loop in ShapeScript instead, with optimizers, datasets and checkpoints.)

What ShapeScript adds at run time is a check on the way in. Swap the weights by mistake, loss(x, y, w2, w1), and the call fails at once, raising a ShapeError that names the parameter, its type and both shapes. shapescript run shows it as the checker shows its errors, with the argument underlined and the parameter it was given for:

$ shapescript run train.py
error: argument `w1` of `loss` should be [D Hd] = (2, 32), got shape (32, 3)
  --> train.py:14:20
   |
14 |     l = loss(x, y, w2, w1)
   |                    ^^ expected [D Hd] = (2, 32), got shape (32, 3)
   |
  ::: classifier.ss:15:36
   |
15 | fn loss(x: [b D], y: [b] Index[C], w1: [D Hd], w2: [Hd C]) -> [] NonNeg:
   |                                    ---------- parameter declared here
   |
   = note: run with `SHAPESCRIPT_TRACEBACK=1` to see the Python traceback

Making mistakes on purpose

The real point of the checker is the mistakes it catches before anything runs. Here are three that happen in real code.

Applying softmax before the loss. cross_entropy takes logits, and applies log_softmax itself. Passing it probabilities still trains, badly, and nothing in PyTorch complains. ShapeScript knows what softmax returns:

dim D = 2
dim Hd = 32
dim C = 3

fn forward(x: [b D], w1: [D Hd], w2: [Hd C]) -> [b C] Logit over C:
    return x @ w1 |> relu |> matmul(w2)

fn loss(x: [b D], y: [b] Index[C], w1: [D Hd], w2: [Hd C]) -> [] NonNeg:
    p = softmax(forward(x, w1, w2), over=C)
    return cross_entropy(p, y, over=C)

$ shapescript check softmax_twice.ss
error[E0501]: expected `Logit over C`, got `Prob over C`
  --> softmax_twice.ss:10:26
   |
10 |     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`.

Mixing up axes. Suppose you average the scores over the batch to get a “typical” score per class, but write the wrong axis:

dim C = 3

fn mean_scores(z: [b C]) -> [C]:
    return mean(z, over=C)

$ shapescript check wrong_axis.ss
error[E0401]: mismatched shape: expected `[C]`, found `[b]`
 --> wrong_axis.ss:4:12
  |
4 |     return mean(z, over=C)
  |            ^^^^^^^^^^^^^^^ the return type is `[C]`
  |
  = note: axis 1: `b` is not known to equal `C` = 3

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

Averaging over C leaves [b], which isn’t [C]. The fix is over=b.

A label from the wrong set. Labels typed Index[C] are positions along an axis of size C. Using them to index an axis of another size is an error, where PyTorch would read the wrong rows or fail with an index error at run time:

dim C = 3
dim K = 5

fn pick(table: [K 4], y: [b] Index[C]) -> [b 4]:
    return table[y]

$ shapescript check wrong_index.ss
error[E0507]: expected `Index[K]`, got `Index[C]`
 --> wrong_index.ss:5:18
  |
5 |     return table[y]
  |            ----- ^ this is `[b] Index[C]`
  |            |
  |            the axis it indexes is `K` = 5
  |
  = note: the sizes differ: `C` = 3, not `K` = 5

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

Exercises

1. Add a bias to the hidden layer: a parameter b1: [Hd] added after x @ w1. Where does the bias go in the signature, and does + need an unsqueeze?
2. Write accuracy(x: [b D], y: [b] Index[C], w1: [D Hd], w2: [Hd C]), returning the fraction of correct predictions as a []. Hint: comparing two tensors with == gives a mask, and to f32 turns a mask into numbers you can average.
3. Predict the type of forward(x, w1, w2) where x: [16 D].

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.

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

Pattern Operations

The spotter can normalize tokens, pool them and score them, but it doesn’t have any tokens yet. It needs two things that are both about moving axes around: cutting each image into patches, a token per patch, and mixing information between the patches. Both are reshaping, and reshaping is where shape bugs breed. PyTorch’s view, permute, transpose and reshape take positions and sizes, and a wrong one often still runs. einops fixed much of this with patterns such as "b n (h d) -> b h n d", but they’re strings, checked only when they run, and the names in one pattern mean nothing to the next line.

ShapeScript builds einops’ patterns into the language. Five operations take one: rearrange, reduce, repeat, einsum and split. This chapter uses them to give the spotter its patches and its first mixing layer.

Cutting images into patches

An image is [b 1 S S]: a batch, one channel, and 16 × 16 pixels. We want [b N P*P]: for each image, N = 16 patches, each a row of the P × P = 16 pixels it covers. Add this to spotter.ss:

## 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 pattern is inputs -> output, written without quotes. The left side describes the operand’s axes by position, one entry per axis, and each entry binds to one axis of img:

• A name already in scope must have that axis’s size. b is img’s batch axis, so it matches.
• A new name is introduced with the size of the axis it binds to, and stays in scope for the rest of the function. c binds the channel axis, of size 1.
• A number must equal the axis’s size.
• A group (gh p1) binds one axis whose size is the product of its parts. At most one part may be unknown, and the checker solves for it.

Here the groups bind the two axes of size S. A new name can get its size from a keyword argument, as in einops, and the size may be a number or a dim: p1=P says each patch is 4 pixels high, so the checker solves gh = 16 / 4 = 4, and the same for gw. Those are the 4 × 4 patches down and across.

The right side says what the result looks like. Its names must come from the left (except in repeat), and the result’s axes are named after them: the result is [b (gh gw) (p1 p2 c)], which is [b 16 16]. That fits the return type [b N P*P], because the checker compares sizes after working them out: gh × gw = 16 = N, and p1 × p2 × c = 16 = P*P.

Sizes that don’t divide are an error. Try 5-pixel patches:

dim S = 16

fn patchify(img: [b 1 S S]) -> [b ?n ?k]:
    return img |> rearrange(b c (gh p1) (gw p2) -> b (gh gw) (p1 p2 c), p1=5, p2=5)

$ shapescript check bad_patches.ss
error[E0402]: can't split this axis into `(gh p1)`
 --> bad_patches.ss:4:33
  |
4 |     return img |> rearrange(b c (gh p1) (gw p2) -> b (gh gw) (p1 p2 c), p1=5, p2=5)
  |                                 ^^^^^^^ this axis is `S` = 16
  |
  = note: 16 isn't divisible by 5
  = note: in stage 1 of this pipeline, whose input is `[b 1 S S]`

error[E0402]: can't split this axis into `(gw p2)`
 --> bad_patches.ss:4:41
  |
4 |     return img |> rearrange(b c (gh p1) (gw p2) -> b (gh gw) (p1 p2 c), p1=5, p2=5)
  |                                         ^^^^^^^ this axis is `S` = 16
  |
  = note: 16 isn't divisible by 5
  = note: in stage 1 of this pipeline, whose input is `[b 1 S S]`

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

One error for each side of the image. In PyTorch, img.unfold(2, 5, 5) would quietly drop the last pixel of each row and column.

Names are positions, not lookups

The checker catches a classic einops mistake: a pattern that names an axis after a different axis of the operand. Patterns bind by position, so this almost always means the axes were written in the wrong order. The tokens the spotter will mix are [b N D]; here is an attempt to turn them into [b D N] that writes the left side in the order wanted on the right:

dim N = 16
dim D = 32

fn swap(x: [b N D]) -> [b D N]:
    return x |> rearrange(b D N -> b N D)

$ shapescript check transposed.ss
error[E0408]: the pattern calls axis 2 `D`, but `D` is axis 3 of `[b N D]`
 --> transposed.ss:5:29
  |
5 |     return x |> rearrange(b D N -> b N D)
  |            -                ^ axis 2 is `N`
  |            |
  |            this is `[b N D]`
  |
  = note: patterns name axes by position
  = help: to reorder the axes, use `rearrange` first
  = note: in stage 1 of this pipeline, whose input is `[b N D]`

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

The left side describes what x is, b N D, and the right side what you want, b D N.

rearrange

rearrange moves, merges and splits axes. Every name on the left must be on the right, so no data is lost. patchify merges, and the fix above moves:

dim N = 16
dim D = 32

fn swap(x: [b N D]) -> [b D N] = x |> rearrange(b N D -> b D N)

fn flatten(img: [b c h w]) -> [b (c h w)] = rearrange(img, b c h w -> b (c h w))

fn merge_heads(x: [b H n Dh]) -> [b n (H Dh)] = rearrange(x, b H n Dh -> b n (H Dh))

A 1 on the right adds an axis of size 1, as unsqueeze would. That’s how to give a single image a batch axis:

dim S = 16

fn as_batch(img: [1 S S]) -> [1 1 S S] = rearrange(img, c h w -> 1 c h w)

reduce

reduce removes the axes left out of the right side, combining them with sum, mean, max, min or prod. Chapter 3’s nonblank used one: reduce(img, b c h w -> b, max) is the brightest pixel of each image. Groups work here too. The mean brightness of each patch, which tells you where the ink is, is patchify with a reduction instead of a merge:

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

## The mean brightness of each patch.
fn patch_brightness(img: [b 1 S S]) -> [b N]:
    return img |> reduce(b c (gh p1) (gw p2) -> b (gh gw), mean, p1=P, p2=P)

## Halves each side of an image, keeping the brightest pixel of each 2×2 block.
fn shrink(img: [b 1 S S]) -> [b 1 S/2 S/2]:
    return img |> reduce(b c (h p1) (w p2) -> b c h w, max, p1=2, p2=2)

For a single named axis, the reductions are also functions that take the axis with over=: the spotter’s pool, mean(x, over=n), is reduce(x, b n D -> b D, mean). sum(x) without over= sums everything to a [].

repeat

repeat adds axes, which may be new names given a size as a keyword. Doubling each pixel is the reverse of shrink:

dim S = 16

fn grow(img: [b 1 S S]) -> [b 1 2*S 2*S]:
    return img |> repeat(b c h w -> b c (h p1) (w p2), p1=2, p2=2)

fn copies(img: [b 1 S S]) -> [b 4 S S] = repeat(img, b c h w -> b (c r) h w, r=4)

Mixing the patches with einsum

Now the spotter’s first mixing layer. After patchify and the step that turns each patch into a token (the next chapter writes it), an image is [b N D]: a token of width D for each of the N patches. A token-mixing MLP runs along the patch axis: for each channel separately, it combines the values of all 16 patches, so one patch learns what’s in the others. That’s an MLP whose weights meet the axis N, with a hidden width of its own. Add a dim for it (why it’s 64 comes up below):

## Hidden width of the MLP that mixes patches.
dim Dt = 64

With rearrange, you’d move N to the end, multiply by w1: [N Dt], apply gelu, multiply by w2: [Dt N], and move it back. einsum says which axes meet directly, so there is nothing to move. It takes several operands, written before a :, and one pattern side per operand. Names shared by several operands must have one size, and names left out of the output are summed over:

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

The first einsum sums over N, the axis that the token and w1 share, and gives [b Dt D]: for each channel, Dt hidden values, each a combination of all the patches. The second sums over Dt and gives back [b N D], which is added to x, so the block is residual. norm is the function from Chapter 4.

einsum is the general product: matrix products, batched dot products and outer products are all patterns. Its operands can be named anything, but the pattern’s names say which axes meet. If you write the output in the wrong order, the return type catches it. Say the first half of the MLP is a function of its own:

dim N = 16
dim D = 32
dim Dt = 64

## The hidden values of the token-mixing MLP.
fn hidden(x: [b N D], w1: [N Dt]) -> [b Dt D]:
    return einsum(x, w1: b N D, N Dt -> b D Dt) |> gelu

$ shapescript check bad_mix.ss
error[E0401]: mismatched shape: expected `[b Dt D]`, found `[b D Dt] [-0.1701, inf)`
 --> bad_mix.ss:7:12
  |
7 |     return einsum(x, w1: b N D, N Dt -> b D Dt) |> gelu
  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the return type is `[b Dt D]`
  |
  = note: axis 2: `D` = 32, not `Dt` = 64
  = note: axis 3: `Dt` = 64, not `D` = 32

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

This is why the spotter’s Dt is 64 and not 32, like D. Dims are compared by their values (Chapter 3), so with Dt = 32 the slip would check: [b D Dt] would be [b 32 32] either way. Give different roles different sizes, and the checker can tell them apart. In mix_patches itself, the same slip in the second einsum, -> b D N, would do no harm: x + ... adds a [b N D] to a [b D N], and the next chapter shows that the axes then line up by name, so the sum is right anyway.

split

split cuts one axis into several results. The pattern’s left side has a number that’s missing from the right; that’s how many results there are, and you bind them to that many names. The spotter has no use for it yet, but it could compare the top half of each image with the bottom half:

dim S = 16

## The top half of each image, minus the bottom half.
fn halves(img: [b 1 S S]) -> [b 1 S/2 S]:
    top, bottom = img |> split(b c (2 h) w -> b c h w)
    return top - bottom

(2 h) on an axis of size 16: the checker works out h = 8, makes two results of [b c h w], and checks that two names receive them. In a transformer, split is how one matrix product makes the queries, keys and values: q, k, v = x @ w_qkv |> split(b n (3 H Dh) -> b H n Dh).

Operations along an axis

Many operations work along one axis, and take it by name with over=. The spotter’s class scores are [b C], and most of what you’d do with them is one of these:

dim S = 16
dim C = 4

fn probs(scores: [b C]) -> [b C] Prob over C = softmax(scores, over=C)

fn best(scores: [b C]) -> [b] Index[C] = argmax(scores, over=C)

## The two likeliest shapes of each image.
fn best_two(scores: [b C]) -> [b 2] Index[C]:
    values, ids = topk(scores, 2, over=C)
    return ids

## Two batches of images, one after the other.
fn together(a: [b 1 S S], c: [m 1 S S]) -> [b+m 1 S S] = cat(a, c, over=b)

fn running_total(x: [b n]) -> [b n] = cumsum(x, over=n)

cat adds up the sizes along its axis, so the result is [b+m 1 S S]. stack(a, c, over=k) puts its arguments along a new first axis named k. topk gives two results, the values and their positions.

Indexing and slicing

Indexing removes an axis, and slicing shortens one:

dim S = 16
dim D = 32

fn first_image(img: [b 1 S S]) -> [1 S S] = img[0]

fn first_patch(x: [b n D]) -> [b D] = x[:, 0]

fn drop_last(x: [b n D]) -> [b n-1 D] = x[:, :-1]

fn shifted(x: [b n D]) -> [b n-1 D] = x[:, 1:]

The checker works out a slice’s size when it can prove the slice fits the axis: :-1 and 1: on an axis n both have n-1 elements. A slice that might run past the end, such as x[:, 0:3] on an axis whose size is generic, gets an unknown size.

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
## Hidden width of the MLP that mixes patches.
dim Dt = 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] = 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)

Exercises

1. Change the patches to 8 × 8 pixels. What do you change in spotter.ss, what are N and P*P now, and does patchify need to change?

2. Write unpatchify(x: [b N P*P]) -> Images, which puts the patches back together, so that unpatchify(patchify(img)) is img. Hint: from one axis of size 16, the checker can’t tell how many patches go down and how many across, so give one of them as a keyword. A keyword size can be a number or a dim, but not S / P, which is a float.

3. This first try at the token-mixing MLP uses rearrange and @ instead of einsum. It has a mistake. Find it with the checker, and fix it:

   dim N = 16
   dim D = 32
   dim Dt = 64
   
   fn mix_patches(x: [b N D], w1: [N Dt], w2: [Dt N]) -> [b N D]:
       h = x |> rearrange(b N D -> b D N) |> matmul(w1) |> gelu |> matmul(w2)
       return x + rearrange(h, b N D -> b D N)
   

4. Write ink_map(img: [b 1 S S]), the brightest pixel of each patch, as a 4 × 4 grid per image. What’s its return type?

Broadcasting by Name, @ and Pipes

The spotter has most of its parts now: patchify cuts an image into patches, norm normalizes tokens, mix_patches mixes the patches, and head turns tokens into class scores. What’s missing is the glue: a matrix product that turns each patch into a token, a second MLP that mixes the channels of each token, and a forward function that runs the whole thing in order.

That glue is operators: how +, * and friends line up tensors of different shapes, how @ multiplies matrices, and how |> strings operations together. This chapter covers all three, and ends with a spotter that goes from pixels to scores.

Broadcasting by name

Look again at norm, from Chapter 4:

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

mu is [b n], one mean per token, and x is [b n D]. In PyTorch, x - mu would fail, or worse, quietly line up the wrong axes: you’d write mean(dim=-1, keepdim=True), or mu[..., None], and getting the Nones wrong is a silent bug when sizes happen to match. Here x - mu just works, because ShapeScript lines axes up by name whenever it can. mu’s axes are b and n, x has axes called b and n, so that’s where they go, and D is broadcast.

A few more that need no unsqueeze:

dim D = 64

fn add_mask(scores: [b H n m], mask: [n m]) -> [b H n m] = scores + mask

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

fn scale_rows(x: [b n D], w: [b n]) -> [b n D] = x * w

## Each image minus its own mean brightness.
fn center(img: [b c h w]) -> [b c h w] = img - reduce(img, b c h w -> b, mean)

The compiled Python inserts the right Nones.

The rule has two cases:

1. By name. If every axis of one operand has a name, and all those names are axes of the other operand, they line up by name. The result has the other operand’s axes, in its order. (If both qualify, as with [n m] + [m n], the left operand’s order wins.)
2. From the right. Otherwise the axes line up from the right, as in NumPy. Each pair must have the same size or one of them size 1, and two named axes in a pair must have the same name.

The second case is what catches bugs. Say you want to give each patch a learned position vector, so the spotter knows where each patch was, and you write the parameter with a fresh axis name:

dim N = 16
dim D = 32

fn add_positions(x: [b N D], pos: [k D]) -> [b N D] = x + pos

$ shapescript check positions.ss
error[E0404]: can't line up the axes of `[b N D]` and `[k D]`
 --> positions.ss:4:55
  |
4 | fn add_positions(x: [b N D], pos: [k D]) -> [b N D] = x + pos
  |                                                       -^^^---
  |                                                       |   |
  |                                                       |   this is `[k 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 `N` meets `k`
  = help: give them the same names with `rearrange` or `as`

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

k isn’t an axis of x, so the axes line up from the right, and N meets k. Even if a caller passes 16 rows, they’re different axes as far as the function knows. The fix is to say what you mean, pos: [N D]: one vector per patch. If two axes really are the same but have different names, say so: Chapter 7 shows as, and Chapter 5 rearrange.

Unnamed axes, such as the 1 in [b 1 S S] or the 3 of an RGB image, line up from the right only:

## Adds a per-channel offset to an RGB image.
fn shift(img: [b 3 h w], offset: [3 1 1]) -> [b 3 h w] = img + offset

Matrix multiplication

The first step of the spotter turns each patch, 16 pixels, into a token of D numbers. That’s a matrix product. x @ y and matmul(x, y) work like torch.matmul: the last axis of x meets the second-to-last of y (or its only axis), and the axes before those broadcast by the rule above. Add this to spotter.ss:

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

[b N P*P] @ [P*P D] is [b N D]: the pixels of each patch meet the rows of w_e, and b and N come along. Batched products work the same way, with the leading axes broadcast:

fn per_head(q: [b H n Dh], k: [b H Dh m]) -> [b H n m] = q @ k

The contracted axes must have equal sizes, but not equal names. That’s deliberate: in x @ w, the meeting axes usually have different roles. In embed it’s “the pixels of a patch” meeting “the inputs of the layer”; they’re both P*P here, but they could as well be Din and P*P. A size mismatch is still an error. Write the weights the other way round, as nn.Linear stores them:

dim P = 4
dim D = 32

fn embed(patches: [b n P*P], w_e: [D P*P]) -> [b n D] = patches @ w_e

$ shapescript check bad_embed.ss
error[E0401]: can't multiply `[b n P*P]` by `[D P*P]`
 --> bad_embed.ss:4:57
  |
4 | fn embed(patches: [b n P*P], w_e: [D P*P]) -> [b n D] = patches @ w_e
  |                                                         -------^^^---
  |                                                         |         |
  |                                                         |         this is `[D P*P]`
  |                                                         this is `[b n P*P]`
  |
  = note: the axes `@` contracts differ: `P*P` = 16, not `D` = 32

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

Pipes

The second MLP of a mixer block works on each token on its own, mixing its D channels: normalize, widen to Dc, gelu, narrow back to D, and add the input. Written as nested calls, that’s add(x, matmul(gelu(matmul(norm(x), w1)), w2)), which you read inside out. A pipe reads in the order it runs. Add this:

## Hidden width of the MLP that mixes channels.
dim Dc = 64

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

a |> f(args) means f(a, args), and a |> f means f(a). A few rules:

• _ puts the value somewhere other than first. |> add(x, _) is add(x, value), the residual connection. A stage may have one _.
• A line that starts with |> continues the line above, so a pipeline can be written one stage per line, as below.
• |> binds more loosely than everything else. x @ w |> gelu is (x @ w) |> gelu. To apply as or to to a whole pipeline, use parentheses: (x |> f) as T.

Every built-in takes the tensor first, so any of them can be a stage, and so can your own functions: norm above, and, in the next listing, embed, mix_patches, mix_channels and head. That’s the whole spotter, one stage per line:

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

Read it top to bottom: patches become tokens, the tokens mix with each other, then within themselves, and are normalized, averaged and scored. (A real mixer has several blocks; Chapter 9 adds a loop.) The parameter list continues on a second line, since it’s inside parentheses.

When a stage fails, the message says which one, and what its input was. Here mix_channels uses the first weight twice:

dim N = 16
dim D = 32
dim Dc = 64

fn mix_channels(x: [b N D], w1: [D Dc], w2: [Dc D]) -> [b N D]:
    return x
        |> layer_norm(over=D)
        |> matmul(w1)
        |> gelu
        |> matmul(w1)
        |> add(x, _)

$ shapescript check bad_pipe.ss
error[E0401]: can't multiply `[b N Dc] [-0.1701, inf)` by `[D Dc]`
  --> bad_pipe.ss:10:12
   |
 6 |       return x
   |  ____________-
 7 | |         |> layer_norm(over=D)
 8 | |         |> matmul(w1)
 9 | |         |> gelu
   | |_______________- this is `[b N Dc] [-0.1701, inf)`
10 |           |> matmul(w1)
   |              ^^^^^^^--^
   |                     |
   |                     this is `[D Dc]`
   |
   = note: the axes `@` contracts differ: `Dc` = 64, not `D` = 32
   = note: in stage 4 of this pipeline, whose input is `[b N Dc] [-0.1701, inf)`

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

The second matrix should have been w2. (The [-0.1701, inf) after the shape is the range of values gelu can give, which the checker tracks; Chapter 8 explains why.)

In the editor, each line of a pipeline ends with its shape, so you can see this before you run anything (Chapter 13).

The spotter so far

spotter.ss now goes from an image to four class scores. It can’t run anything useful yet, since it has no trained weights; the next chapter calls it from Python for the first time.

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] = 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, w_e: [P*P D], tw1: [N Dt], tw2: [Dt N],
           cw1: [D Dc], cw2: [Dc D], w_h: [D C]) -> [b C]:
    return img
        |> embed(w_e)
        |> mix_patches(tw1, tw2)
        |> mix_channels(cw1, cw2)
        |> norm
        |> head(w_h)

Exercises

1. Give embed a bias, b_e: [D], added after the matrix product. Does the + need an unsqueeze? Then add a learned position vector per patch, pos: [N D], so the spotter knows where each patch came from.
2. Real layer norms have a learned gain and shift. Give norm two more parameters, g: [D] and beta: [D], and apply them to the result. Which broadcasting rule lines them up?
3. Rewrite head as a pipeline that starts with x.
4. mix_patches uses einsum so that the weights meet the patch axis N, not the last axis. With the weights stored the other way round, w1: [Dt N] and w2: [N Dt], it’s w1 @ norm(x): a [Dt N] times a [b N D]. Write this version of mix_patches as a pipeline that starts with x. Hint: _.

Assertions with as and Runtime Checks

The spotter has a whole forward pass now: patches, a token per patch, a mixer block, and class scores. Inside spotter.ss, the checker has proved every shape. But the images and weights come from Python, and the checker can’t see Python. This chapter runs the spotter from Python for the first time and covers what protects it there: the checks that run when the program does, and as, which tells the checker something it can’t prove.

Running the spotter

First we need images. Make shapes.py next to spotter.ss. It draws squares, disks, pluses and crosses of random sizes in random places, on a noisy background, and returns them as bytes, the way images are usually stored:

import torch

CLASSES = ["square", "disk", "plus", "cross"]


def draw(n, seed=0):
    """`n` 16×16 grayscale images of shapes, as bytes, and their labels."""
    g = torch.Generator().manual_seed(seed)
    y = torch.randint(0, 4, (n,), generator=g)
    r = torch.randint(2, 7, (n, 1, 1), generator=g).float()  # half the size
    cy = torch.randint(5, 11, (n, 1, 1), generator=g).float()
    cx = torch.randint(5, 11, (n, 1, 1), generator=g).float()
    i = torch.arange(16).float()
    dy, dx = i.view(1, 16, 1) - cy, i.view(1, 1, 16) - cx
    shapes = torch.stack([
        (torch.maximum(dy.abs(), dx.abs()) - r).abs() < 0.6,         # square
        dy**2 + dx**2 <= r**2,                                        # disk
        ((dy.abs() < 1) | (dx.abs() < 1)) & (dy.abs() <= r) & (dx.abs() <= r),  # plus
        ((dy - dx).abs() < 1) | ((dy + dx).abs() < 1),                # cross
    ])
    ink = shapes[y, torch.arange(n)]
    ink[y == 3] &= (dy.abs() <= r)[y == 3]
    pixels = ink.float() * 140 + 110 * torch.rand(n, 16, 16, generator=g)
    return pixels.to(torch.uint8).unsqueeze(1), y

There’s nothing to learn from that file except what it returns: a (n, 1, 16, 16) tensor of bytes, and n labels from 0 to 3, one per class in CLASSES.

Now run.py, which pushes four images through the spotter with random weights:

import torch
from shapes import CLASSES, draw
from spotter import forward

img, y = draw(4)
img = img / 255  # bytes to numbers in [0, 1]

g = torch.Generator().manual_seed(0)
w = [torch.randn(16, 32, generator=g), torch.randn(16, 64, generator=g),
     torch.randn(64, 16, generator=g), torch.randn(32, 64, generator=g),
     torch.randn(64, 32, generator=g), torch.randn(32, 4, generator=g)]

scores = forward(img, *w)
print(scores.shape)
for s, label in zip(scores.argmax(-1), y):
    print("guessed", CLASSES[s], "for a", CLASSES[label])

img / 255 turns the bytes into numbers in [0, 1], since forward takes f32 images; Chapter 8 lets the spotter take the bytes itself. The weights are in the order of forward’s parameters: w_e, tw1, tw2, cw1, cw2 and w_h. Run it:

$ shapescript run run.py
torch.Size([4, 4])
guessed disk for a square
guessed disk for a cross
guessed disk for a disk
guessed disk for a square

The spotter runs: four images in, four rows of four scores out. The guesses are no better than chance, because the weights are random: this one happens to think everything is a disk. Chapter 10 trains them.

Runtime checks

forward is ordinary Python now, and Python will call it with anything. So every ShapeScript function checks its arguments when it’s called, whether from Python or from other ShapeScript: the number of axes, each fixed size, that the generic sizes agree across the arguments, and that each holds the right kind of values: floats, integers or booleans, as its dtype says. Here are two mistakes worth making on purpose. A batch of 28×28 images, the size of MNIST’s digits:

forward(torch.rand(4, 1, 28, 28), *w)

and the patch-mixing weights in the wrong order:

w_e, tw1, tw2, cw1, cw2, w_h = w
forward(img, w_e, tw2, tw1, cw1, cw2, w_h)

Each fails at once, at the call, naming the parameter, its type, and both shapes:

ShapeError: spotter.ss:54:12: argument `img` of `forward` should be [b 1 S S] = (?, 1, 16, 16), got shape (4, 1, 28, 28)
ShapeError: spotter.ss:54:39: argument `tw1` of `forward` should be [N Dt] = (16, 64), got shape (64, 16)

(The messages give the file’s full path; it’s shortened here.) Without the checks, both mistakes would still fail, but deep inside: in rearrange, with a message about einops’ pattern, and in an einsum, with one about operand sizes. And if Dt were 16, so that both weights were [16 16], the swap wouldn’t fail at all.

ShapeError is a ValueError, from the shapescript package, as are the other errors in this chapter. A traceback shows the lines of the .ss file, not of generated Python.

as claims a type

The checks on the way in protect a function’s signature. Sometimes you want to go further, and make a claim the checker can’t prove. x as T claims that x already has the type T. The checker believes it from then on, and the compiled program checks it when it runs.

Say the images come from a folder of files, and Python loads them as [b 1 h w], whatever their size. They’re supposed to be 16×16. Add to spotter.ss:

## Images of any size from Python, which must be S × S.
fn from_python(img: [b 1 h w]) -> Images = img as Images

h and w are generic, so they might be 16 at run time, and the claim is allowed. Without the as, returning img would be an error, because h isn’t known to equal S. With it, the rest of the program sees Images, and passes it to forward without complaint.

Unknown sizes

Chapter 3’s nonblank keeps the images that have any ink, so it returns [?k 1 S S]: how many it keeps is only known when it runs. Often you know more. If every image in a batch should have ink, the number kept is the number of images, and you can say so:

## The images, which must all have ink.
fn inked(img: Images) -> Images = nonblank(img) as [b 1 S S]

?k becomes b, the batch size of img. If an image is blank after all, the claim is false, and the program stops there instead of carrying on with three images and four labels.

as is the only way to make a claim the checker can’t prove, and it always compiles to a runtime check. So every place a program relies on an assumption is a visible as, and you can find them all with a search.

What as can claim

An as must be possible. A claim that contradicts what the checker knows is an error, because it can only fail:

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

fn upscale(img: Images) -> [b 1 28 28]:
    return img as [b 1 28 28]

$ shapescript check upscale.ss
error[E0410]: `[b 1 S S]` can never be `[b 1 28 28]`
 --> upscale.ss:5:12
  |
5 |     return img as [b 1 28 28]
  |            ---^^^^^^^^^^^^^^^
  |            |
  |            this is `[b 1 S S]`
  |
  = note: axis 3: `S` = 16, not 28

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

as never changes a value, so it can’t resize an image. Symbolic sizes are fine, since they might be equal at run time, as h and S were in from_python. That makes as the way to say two generic axes are the same. Blending two batches of images, pixel by pixel, needs them to be equally long:

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

## The average of two batches of images, which must be equally long.
fn blend(a: Images, c: [m 1 S S]) -> Images = (a + (c as [b 1 S S])) / 2

Without the as, a + c is error E0404 (Chapter 6), because b and m are different axes. With it, the program checks at run time that c has b images.

as also gives a type to a value the checker knows nothing about, such as one from Python (Chapter 10). The parts of the type you don’t write are kept: x as [b C] keeps x’s domain, and x as Unit keeps its shape.

Claiming values

A type can also say what values a tensor holds: its domain, which Chapter 8 is about. Unit means values in [0, 1], which is what forward expects of its pixels. run.py divided by 255 to get there; forgetting to is a classic bug, which trains on numbers 255 times too large. Claim the range on the way in:

## The scores of images whose pixels must be in [0, 1].
fn spot(img: Images, w_e: [P*P D], tw1: [N Dt], tw2: [Dt N],
        cw1: [D Dc], cw2: [Dc D], w_h: [D C]) -> [b C]:
    return forward(img as Unit, w_e, tw1, tw2, cw1, cw2, w_h)

When claims fail

Each claim is checked where it’s made. Called from Python with the wrong data, every one of these raises an error that names the .ss file, line and column, and what was expected:

from_python(torch.rand(4, 1, 28, 28))  # not 16×16
blank = img.clone()
blank[2] = 0
inked(blank)                           # image 2 has no ink
spot(img * 255, *w)                    # bytes, not [0, 1]

ShapeError: spotter.ss:63:44: expected [b 1 S S] = (4, 1, 16, 16), got shape (4, 1, 28, 28)
ShapeError: spotter.ss:66:35: expected [b 1 S S] = (4, 1, 16, 16), got shape (3, 1, 16, 16)
DomainError: spotter.ss:71:20: expected Unit (values in [0, 1]), got values in [0, 249]

The second message is the claim of inked failing: nonblank kept 3 images, not b = 4. The third is a DomainError, a check of the values rather than the shape.

as versus to

as never changes a value; it only states what the value is. to, which Chapter 8 covers, converts. img as Unit checks that the pixels are in [0, 1]; img to Unit, for bytes, divides them by 255. Keep the two apart:

Turning checks off

The checks are cheap, but not free, so shapescript build and the Python import hook let you choose:

debug means “on, unless Python runs with -O”. So python -O would let spot(img * 255, *w) through, and the shape checks would still run. From Python, shapescript.install(checks="always", value_checks="off") sets them for imported files.

On a GPU, value checks don’t wait for the result: they use torch._assert_async, so a failure shows up at a later CUDA call as a device-side assert. To see which check failed, run the same inputs on the CPU.

Seeing the generated code

shapescript build writes the Python that a file compiles to. It’s worth reading once, to see that there’s no magic:

shapescript build spotter.ss -o build/spotter.py

The function inked becomes:

def inked(img):
    """The images, which must all have ink."""
    _ss_dev = _ss_rt.device()
    _ss_rt.check_rank(img, 4, "inked", "img", "[b 1 S S]", "spotter.ss:66:10")
    _ss_rt.check_dtype(img, _ss_torch.float32, "f32", "inked", "img", "[b 1 S S]", "spotter.ss:66:10")
    _ss_rt.check_device(img, _ss_dev, "inked", "img", "spotter.ss:66:10")
    _ss_d_b = img.shape[0]
    _ss_rt.check_sizes(img, (None, 1, S, S), "inked", "img", "[b 1 S S]", "spotter.ss:66:10")
    return _ss_check(nonblank(img), (_ss_d_b, 1, S, S), "spotter.ss:66:35", "[b 1 S S]")

The first four checks are on the way in: the number of axes, the kind of values (floats here, since the type gives no dtype and so means f32; bytes or booleans would be an error), the device (Chapter 16), and the sizes, with None for the generic b, whose size is read from the argument. The as is the _ss_check on the last line, which compares what nonblank returned with that b.

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
## 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] = 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, w_e: [P*P D], tw1: [N Dt], tw2: [Dt N],
           cw1: [D Dc], cw2: [Dc D], w_h: [D C]) -> [b C]:
    return embed(img, w_e)
        |> mix_patches(tw1, tw2)
        |> mix_channels(cw1, cw2)
        |> norm
        |> head(w_h)

## Images of any size from Python, which must be S × S.
fn from_python(img: [b 1 h w]) -> Images = img as Images

## The images, which must all have ink.
fn inked(img: Images) -> Images = nonblank(img) as [b 1 S S]

## The scores of images whose pixels must be in [0, 1].
fn spot(img: Images, w_e: [P*P D], tw1: [N Dt], tw2: [Dt N],
        cw1: [D Dc], cw2: [Dc D], w_h: [D C]) -> [b C]:
    return forward(img as Unit, w_e, tw1, tw2, cw1, cw2, w_h)

Exercises

1. Write crop(img: [b 1 h w]) -> Images, which keeps the top-left S×S corner of each image, for photos larger than 16×16. A slice is img[:, :, :S, :S]; why does it need an as?

2. Make blend check without as, by changing only its signature. When is that the better fix?

   dim S = 16
   dim Images = [b 1 S S]
   
   fn blend(a: Images, c: [m 1 S S]) -> Images = (a + c) / 2
   

3. A colleague puts x = x as [b N D] after every stage of forward, “to be safe”. What does that cost, and what would you do instead?

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)
   

Control Flow and Iterator Chains

The spotter works, but it’s shallow: one block that mixes the patches, one that mixes the channels. Real mixers stack several. And when it gets an image wrong, a single accuracy figure doesn’t say which shapes it confuses. This chapter fixes both. Stacking blocks needs a for loop, whose steps run one after another. Measuring each class needs an iterator chain, which computes something for every element along an axis, all at once. Along the way we’ll meet if.

Stacking blocks with a loop

for i in a..b: runs its block for each whole number from a up to, but not including, b, like Python’s range(a, b).

To stack L blocks, each block needs its own weights. Rather than eight weight parameters for two blocks, and more for three, give each block weight a leading axis L: tw1: [L N Dt] holds the patch-mixing weights of every block, and tw1[l], indexing that axis, is block l’s [N Dt]. Here is the new forward, which replaces Chapter 8’s. Add L to the dims, and change forward, loss and predict in spotter.ss:

## Number of mixer blocks.
dim L = 2

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

forward is no longer one pipeline: it embeds the patches, runs the blocks in turn, then normalizes and classifies. The ends of a range are host numbers: ints, dims or sizes, never tensors, since a loop counts on the host. l is an int, and tw1[l] is a [N Dt], which is what mix_patches takes. There’s no while, break or continue; a return inside a loop returns from the function.

loss and predict take the stacked weights too, and pass them on:

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

A deeper spotter is now one line: dim L = 4. The checker then expects every block weight with 4 in front, and so does the runtime check on the way in from Python.

What a loop may change

A loop’s block may give a new value to a variable from outside it, as x = x |> ... does. The new value must have the same shape, axis names and dtype, so that the code after the loop, and the next time round, work however many times the loop ran. Suppose a stray |> pool, which belongs after the loop, ends up inside it:

fn blocks(x: [b N D], tw1: [L N Dt], tw2: [L Dt N]) -> [b N D]:
    for l in 0..L:
        x = x |> mix_patches(tw1[l], tw2[l]) |> pool
    return x

$ shapescript check loop_bad.ss
error[E0801]: `x` can't change type in a loop
 --> loop_bad.ss:3:13
  |
1 | fn blocks(x: [b N D], tw1: [L N Dt], tw2: [L Dt N]) -> [b N D]:
  |           - `x` is `[b N D]` here
2 |     for l in 0..L:
3 |         x = x |> mix_patches(tw1[l], tw2[l]) |> pool
  |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ this is `[b D]`, but `x` is `[b N D]`
  |
  = help: a value from outside a loop must have the same type on every iteration; give this one a new name

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

The first time round, x would become a [b D], and the second block can’t take that. To compute something of another type, give it a new name. Names first bound inside a block are gone after it.

Loops also run over a dataset’s batches, for batch in data |> batches(64):, which Chapter 12 covers, and over Python values (Chapter 10).

Conditions

if takes a bool or a [] mask, and else: goes on its own line. There’s no elif; nest an if in the else block instead. Say you want to try other activations in the channel MLP, chosen by a number from Python:

dim Dc = 64

fn activation(h: [b n Dc], kind: int) -> [b n Dc]:
    if kind == 0:
        return gelu(h)
    else:
        if kind == 1:
            return relu(h)
        else:
            return tanh(h)

## Shrinks `h` if any of it is large.
fn tame(h: [b n Dc]) -> [b n Dc]:
    if max(abs(h)) > 10:
        h = h * 0.1
    return h

kind == 0 compares two host numbers, a bool. max(abs(h)) > 10 is a [] mask, a value on the device: reading it to decide which way to go makes the program wait for the device, as if tensor.item() > 10 would in PyTorch. A mask with axes isn’t true or false; reduce it with any or all first.

Measuring the spotter with iterator chains

Comparing predictions with labels gives the overall accuracy: mean(pred == y to f32). What we want is one number per class: of the images that really are squares, how many did the spotter call squares? In PyTorch you’d write a loop over the classes, or some careful indexing. In ShapeScript it’s an iterator chain: you write the computation for one class, and say what to iterate.

## 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) / (sum(y == c) to f32))

C |> map(c => e) builds a tensor whose first axis is C, computing e for every index c along it. c => e is a lambda: a parameter, =>, and what to compute from it. Inside, c is an Index[C], a class, so pred == c is a mask of the images predicted as that class. The result has the iterated axis in front, still named C.

The parentheses matter: to binds more tightly than == and and, so sum(pred == c and y == c) to f32 without them would convert c, not the count.

The checker has a warning:

$ shapescript check accuracy.ss
warning[W0501]: `sum(y == c) to f32` can be exactly 0
 --> accuracy.ss:4:63
  |
4 |     return C |> map(c => (sum(pred == c and y == c) to f32) / (sum(y == c) to f32))
  |                                                               ^^^^^^^^^^^^^^^^^^^^ 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`
  = note: in stage 1 of this pipeline, whose input is the indices of `C`

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

It’s right. A batch of 64 random shapes can easily have no crosses in it, and then this class’s accuracy is 0 / 0, a NaN, which would spoil any average you took of it. Divide by at least 1, and an absent class counts as 0:

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

A chain can go inside a chain. The confusion matrix counts, for each true class and each predicted class, how many images had that pair. Its diagonal holds the right answers, and anything off it is a mistake, by kind:

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

A sum of a mask counts its true elements, so the counts are whole numbers, i64. Chapter 10 trains the spotter and prints both; the confusion matrix shows that small squares and small pluses are the shapes it mixes up.

Stages

A chain can start with a size (C, n, 3), a range (0..n) or a tensor, and has four stages:

• map(f): f of every element, stacked along the iterated axis.
• filter(f): the elements for which f is true. How many there are is only known when the program runs, so their axis is an unknown size, ?k, as Chapter 3 described.
• zip(other): pairs of elements from two iterations of the same length, for a lambda with two parameters, (a, b) => ….
• enumerate: pairs of each element’s index and the element.

Starting from a tensor, over= says which axis to iterate, and the elements are the slices along it:

dim C = 4
dim D = 32

## The true class of each image the spotter got wrong.
fn missed(pred: [b] Index[C], y: [b] Index[C]) -> [?k] Index[C]:
    return y
        |> zip(pred, over=b)
        |> filter((t, p) => t != p)
        |> map((t, p) => t)

## Each token's channels, scaled by the token's position.
fn by_position(x: [n D]) -> [n D] = x |> enumerate(over=n) |> map((i, row) => row * (i to f32))

fn smoothed(x: [n D]) -> [n D] = x |> map(gelu, over=n)

The last one gives map a function by name instead of a lambda, as a pipe stage would be: map(gelu) is map(v => gelu(v)), and map(to f32) is map(v => v to f32). In the editor, each line of a chain shows its shape at the end, as for any pipe.

How chains run

A lambda that’s elementwise, like by_position’s, becomes ordinary tensor operations on whole tensors: here one multiplication of x by the positions, as a column. Anything else, like the sums in confusion, becomes torch.vmap, which autograd and torch.export both understand. Either way, no Python loop runs: the 16 entries of the confusion matrix are computed at once, not one after another.

Elements can’t depend on each other

Because the elements are computed at once, one can’t read another. A running sum is a recurrence, which is a loop’s job:

fn running(x: [n]) -> [n]:
    y = n |> map(i => y[i - 1] + x[i])
    return y

$ shapescript check recurrence.ss
error[E2402]: a lambda of an iterator stage can't read `y`, the result it builds
 --> recurrence.ss:2:23
  |
2 |     y = n |> map(i => y[i - 1] + x[i])
  |                       ^ the result of other elements
  |
  = note: an iterator stage computes every element at once, so no element can depend on another
  = help: to compute each element from the one before, use a `for` loop, which runs in order
  = note: in stage 1 of this pipeline, whose input is the indices of `n`

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

(For this one, there’s a built-in: cumsum(x, over=n).) For the same reason a lambda can’t print, call Python, draw random numbers or change a model.

Loops that should be iterator chains

The other way round, the checker notices a loop whose steps don’t depend on each other, and suggests a chain. Here is pool written as a loop over the patches:

dim D = 32

fn pool_loop(x: [b n D]) -> [b D]:
    acc = x[:, 0] * 0
    for i in 0..n:
        acc = acc + x[:, i]
    return acc / (n as Pos)

$ shapescript check hint.ss
hint[H2401]: the iterations of this loop over `n` don't depend on each other
 --> hint.ss:5:5
  |
5 |     for i in 0..n:
  |     ^^^^^^^^^^^^^ runs one iteration after another
6 |         acc = acc + x[:, i]
  |                     ------- the term added to `acc`
  |
  = note: each iteration adds a term to `acc` that doesn't read what the others added, so every term can be computed at once
  = help: compute every term at once with `map`: `acc = acc + (n |> map(i => x[:, i]) |> sum(over=n))`

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

The mixer blocks are different: block 2 takes block 1’s output, so they really must run in order, and the checker says nothing about that loop.

Host scalars

int, float and bool are host scalars: plain Python values for counters, rates and flags, which never go to the GPU. Numbers in the source, loop variables such as l, and sizes used as values are host scalars. Arithmetic on them gives host scalars, and arithmetic with a tensor gives a tensor.

A tensor never turns into a host scalar by itself, because that would make the program wait for the device. Running one more block per mistake, say, is an error:

dim N = 16
dim D = 32
dim C = 4

fn repeat_n(x: [b N D], n: int) -> [b N D]:
    for i in 0..n:
        x = x * 0.5
    return x

fn bad(x: [b N D], pred: [b] Index[C], y: [b] Index[C]) -> [b N D] = repeat_n(x, sum(pred != y))

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

## 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: [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 |> mix_patches(tw1[l], tw2[l]) |> mix_channels(cw1[l], cw2[l])
    return x |> 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)))

Exercises

1. Add accuracy(pred: [b] Index[C], y: [b] Index[C]) -> [] to the spotter: the fraction of all images predicted right.

2. Write class_counts(y: [b] Index[C]) -> [C] i64, how many images of each class a batch has, with map. Then get the same numbers from confusion instead, without a chain.

3. A smaller spotter could use one set of block weights for every block. Write forward_shared, which takes tw1: [N Dt], tw2: [Dt N], cw1: [D Dc] and cw2: [Dc D], and applies the same block L times.

4. This function, meant to drop the first patch after each block, doesn’t check. Why?

   dim D = 32
   dim L = 2
   
   fn shrink(x: [b n D]) -> [b n D]:
       for l in 0..L:
           x = x[:, 1:]
       return x
   

Calling Python and extern fn

At the end of Chapter 9 the spotter was complete: forward, loss, predict and a way to measure it. But its weights are still whatever you pass in, and nothing has trained them. Python is good at that part: loading data, running an optimizer, saving files. ShapeScript doesn’t replace Python; it lives inside it.

This chapter goes both ways. First Python calls the spotter, to train it. Then the spotter calls Python, to load the weights it was trained with: untyped, with import python, and typed, with extern fn.

Training the spotter from Python

Chapter 7 wrote shapes.py, which draws images of shapes with their labels. Create train.py next to it and spotter.ss:

import torch
from shapes import draw
from spotter import loss, predict, per_class_accuracy, confusion

img, y = draw(2000)                     # training images, as bytes, and labels
test_img, test_y = draw(500, seed=1)    # images the model never trains on

g = torch.Generator().manual_seed(0)
def weights(*shape):
    return (torch.randn(*shape, generator=g) / shape[-2] ** 0.5).requires_grad_()

# In the order the spotter's functions take them: w_e, tw1, tw2, cw1, cw2, w_h.
w = [weights(16, 32), weights(2, 16, 64), weights(2, 64, 16),
     weights(2, 32, 64), weights(2, 64, 32), weights(32, 4)]
opt = torch.optim.Adam(w, lr=3e-3)

for step in range(201):
    batch = torch.randint(0, 2000, (64,), generator=g)
    l = loss(img[batch], y[batch], *w)
    opt.zero_grad()
    l.backward()
    opt.step()
    if step % 50 == 0:
        print(f"step {step}: loss {l.item():.3f}")

with torch.no_grad():
    pred = predict(test_img, *w)
print(f"test accuracy: {(pred == test_y).float().mean().item():.1%}")
print(per_class_accuracy(pred, test_y))
print(confusion(pred, test_y))
torch.save(w, "weights.pt")

Each weight starts as random numbers divided by the square root of the axis it’s multiplied along, so the values that flow through the model start out about the size of the inputs. *w passes the six weights in the order the signatures list them. Run it:

$ shapescript run train.py
step 0: loss 1.606
step 50: loss 0.568
step 100: loss 0.151
step 150: loss 0.083
step 200: loss 0.072
test accuracy: 96.4%
tensor([0.9191, 0.9672, 0.9750, 1.0000])
tensor([[125,   1,  10,   0],
        [  0, 118,   1,   3],
        [  2,   1, 117,   0],
        [  0,   0,   0, 122]])

It takes a few seconds on a laptop. The confusion matrix says where the 3.6% of mistakes are: mostly squares taken for pluses, 10 of them, with a few other slips scattered round. At the smallest size, a square outline and a plus are both a few bright pixels in a 4×4 patch, and the model hasn’t quite learned to tell them apart.

Nothing here is special to ShapeScript: loss is an ordinary Python function that calls torch, so autograd, Adam and torch.save work as usual. The last line saves the six trained weights to weights.pt.

import python

Now the other way round. The weights are in a file, and the spotter could load them itself. import python NAME makes a Python module available:

import python torch

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

torch.load(...) is a Python value: the checker knows nothing about it, not even whether it’s a tensor. Indexing it, w[0], or doing arithmetic on it gives another Python value, and you can pass Python values back to Python freely. But where a shape matters, such as passing it to a parameter with a shape, returning it as a tensor, or reducing it over= an axis, it’s an error until an as gives it a type:

import python torch

dim D = 32
dim C = 4

fn saved_head() -> [D C]:
    return torch.load("weights.pt")[5]

$ shapescript check untyped.ss
error[E0413]: the shape of this Python value is unknown
 --> untyped.ss:7:12
  |
7 |     return torch.load("weights.pt")[5]
  |            ^^^^^^^^^^^^^^^^^^^^^^^^^^^ the return type is `[D C]`
  |
  = help: assert its shape with `as [D C]`

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

After the as, everything is checked statically again: predict gets six typed weights, and the checker follows them through the model as before. That’s the pattern for the boundary with Python: one as where a value comes in, and types from there on.

Each as is also checked when the program runs (Chapter 7). Suppose you widen the tokens to dim D = 48, but forget to retrain, so weights.pt still holds the old weights:

ShapeError: spotter.ss:86:25: expected [P*P D] = (16, 48), got shape (16, 32)

The error is at the first weight that doesn’t fit, on the line that loads it, before any of the model runs.

Python values as host values

Python values also work where the program needs a host value: as an if condition, the end of a loop, or what a for loops over. shapes.py has a list of the class names, which a function can print:

import python shapes

fn show_classes():
    for name in shapes.CLASSES:
        print(name)

import python shapes finds shapes.py as Python’s import would, on the Python path, which includes the folder of the script you run. Here name is a Python value too, and print takes anything.

extern fn: typed Python functions

import python is fine for one-off calls. For a function you call often, give it a type with extern fn, naming the Python function in quotes:

## The complementary error function, $1 - \operatorname{erf}(x)$.
extern fn erfc(x: Real) -> NonNeg = "torch.special.erfc"

dim D = 64

fn tail(x: [b D]) -> [b D] NonNeg = erfc(x)

Calls are checked against the signature like any other, and compile to a call of the Python function, with its module imported. The signature is trusted: what the function returns isn’t checked when it runs, so get it right.

Many built-ins are exactly this. exp, log, sigmoid, gelu and the like are extern fn declarations in the prelude, a file every program sees, and linear, layer_norm, rms_norm, conv2d, embedding, scaled_dot_product_attention and dropout are in a second one for torch.nn.functional. The gelu in mix_patches is one of them: in the editor, go to definition on gelu opens its declaration. A declaration in your file with the same name hides the prelude’s.

What an extern fn signature can say

Since there’s no body to check, an extern fn’s signature can say things a fn’s can’t. The spotter has a use for two of them.

The first is a test of how robust the spotter is. A shape moved one pixel to the right is still the same shape, and a good spotter should say so. torch.roll moves the elements of a tensor along one axis, wrapping the last ones round to the front:

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

• ... stands for any number of other axes, the same in every type of the signature. roll takes any shape and gives back the same. The dtype is part of the type, as anywhere: this roll is for the spotter’s bytes.
• dims: Axis[n] takes an axis name, dims=w, and the Python function gets its position, here dims=3. That’s how torch’s positional dim= arguments become named ones. n is the size of that axis, a new name of the signature’s own. As anywhere in a signature, a dim’s name isn’t new: in spotter.ss, Axis[N] would mean an axis of size N = 16, and dims=w would be an error, whose help says so.
• shifts: Size[K] takes a whole number known before the program runs, which becomes the size K in the types.

Why does nudge take [b c h w] and not Images? Because Images is [b 1 S S], which has two axes named S, and dims=S couldn’t say which. In nudge, the axes have names of their own, and a call with Images matches h and w with the two Ses by position, as Chapter 4 described.

Some torch functions take a list of axes, never one number. torch.flip is one: torch.flip(x, 3) is a TypeError. For them, write Axes[n], which takes one axis name too, and passes a list of its position, dims=[3]:

extern fn flip(input: [...] u8, dims: Axes[n]) -> [...] u8 = "torch.flip"

## Each image, mirrored left to right.
fn mirror(img: [b c h w] u8) -> [b c h w] u8 = flip(img, dims=w)

Whether a function wants Axis or Axes is in its documentation, not its name, and the checker can’t know: the signature is trusted. Get it wrong, and the call fails inside torch. That’s the price of an extern fn, and why it’s worth a short test.

The second is a small preview of an image, made by pooling: each output pixel is the brightest pixel of a 3×3 square, with squares every 2 pixels.

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

Division in the return type rounds down, as in torch’s formulas for convolutions and pooling. For a 16-pixel side, (16 − 3) / 2 + 1 is 6.5 + 1 rounded down: 7. The checker works that out from the signature, so -> [b 1 7 7] checks, and [b 1 8 8] wouldn’t.

Here’s a script that loads the trained spotter, tries it on nudged images, and makes thumbnails:

import torch
from shapes import draw
from spotter import predict_saved, nudge, thumbnail

img, y = draw(500, seed=1)
pred = predict_saved(img)
print(f"accuracy: {(pred == y).float().mean().item():.1%}")
print(f"nudged: {(predict_saved(nudge(img)) == y).float().mean().item():.1%}")
print(thumbnail(img.float()).shape)

$ shapescript run check.py
accuracy: 96.4%
nudged: 93.0%
torch.Size([500, 1, 7, 7])

The same accuracy as at the end of training, so the weights came back intact. Moved by one pixel, the images lose 3.4 points: the spotter has learned a little about where the shapes were in its training images, not only what they look like. An exercise below makes it more robust.

Python calling ShapeScript, again

train.py imported spotter.ss as if it were a Python module. A few more details are worth knowing:

• Every ShapeScript function checks its arguments when it’s called, and on the declared device: a tensor on the CPU passed to a program declared devices cuda[0] raises shapescript.DeviceError naming the argument, instead of failing somewhere deep inside.
• There’s no default device. A Python program declares it once, with shapescript.devices("cuda[0]"), unless a .ss file it imports has a devices line, as spotter.ss does.
• shapescript run sets up the import hook for you. A script run with plain python calls shapescript.install() before its first ShapeScript import.
• shapescript.compile(src, path) returns the Python and the diagnostics of some ShapeScript, without running anything.

The spotter so far

devices cpu

import python torch

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

## 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: [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 |> mix_patches(tw1[l], tw2[l]) |> mix_channels(cw1[l], cw2[l])
    return x |> 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. Make the spotter’s predictions steadier: write predict_both, which takes the same arguments as predict, adds the class scores of each image to those of the image nudged one pixel, and returns the likeliest class of the sum.

2. Write an extern fn for torch.nn.functional.avg_pool2d, which takes the same arguments as max_pool2d but averages each square, and use it for tiny(img: Images) -> [b 1 4 4], which averages each 4×4 patch. What would tiny’s return type be with a kernel of 5 and a stride of 3?

3. A colleague saved 100 test images with torch.save(img, "test.pt"). Fix the error, keeping the signature, and say what happens at run time if the file holds 28×28 images:

   import python torch
   
   dim S = 16
   
   fn test_images() -> [?k 1 S S] u8:
       return torch.load("test.pt")
   

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?

Models and Training

The spotter works, but look at what it took. forward, loss and predict each take six weight tensors, in the right order, and pass them on. Python made the weights, ran the training loop and saved them, and predict_saved loaded them back by position, w[0] to w[5], with one as each. That’s how you’d call ShapeScript from an existing PyTorch codebase, and it works. But ShapeScript can also hold the whole training program, and when it does, it checks the parts PyTorch leaves to discipline:

• State is declared. A model’s tensors are the fields of a model declaration. The spotter’s six weights become one value, m: Spotter, so there’s nothing to keep in order and nothing hides in an attribute.
• Gradients are values. value_and_grad(loss, m, ...) returns the loss and its gradient. There’s no .grad to zero and no no_grad mode to forget.
• Changes are visible. A function that changes a model says so with mut, in its signature and at every call.
• Randomness and modes are arguments. dropout takes a Mode, Train(seed=...) or Eval, so evaluation can’t run with dropout left on.
• The loop is yours. There are no hooks or callbacks: the training loop is a function you write and call, and the standard library’s is one you can read on one screen.

The chapter has three parts, and each changes spotter.ss a little more:

1. Models, mut and Gradients: the spotter’s weights become a model, with gradients and an optimizer.
2. Data, Modes and Logs: a record for one labelled picture, dropout, and what training reports.
3. The Training Loop and Checkpoints: the training loop in ShapeScript, saving and resuming, one loop for several models, and the standard library’s fit.

By the end, Python only provides the pictures. The spotter trains itself, and gets 97% of new pictures right.

Exercises

1. In Chapter 10’s train.py, list everything that Python had to get right that the checker couldn’t see: the order of the six weights, their shapes, and so on. Which of them does a model declaration take over?

Models, mut and Gradients

This section turns the spotter’s six weight tensors into one value, a Spotter, and gives it the two things training needs: a gradient, and an optimizer to follow it.

Models

A model declares the tensors that make up a network’s state. In spotter.ss, below the dims, add:

## The spotter's weights.
model Spotter:
    param w_e: [P*P D]  = normal(std=0.25)
    param tw1: [L N Dt] = normal(std=0.25)
    param tw2: [L Dt N] = normal(std=0.125)
    param cw1: [L D Dc] = normal(std=0.18)
    param cw2: [L Dc D] = normal(std=0.125)
    param w_h: [D C]    = normal(std=0.18)

• A param is a tensor that’s trained. Its type is a shape, as always, but one that uses only numbers and dims, so every Spotter has the same shapes.
• Every field has an initializer: zeros, ones, full(value), normal(mean=0, std=1) or uniform(lo=0, hi=1). The standard deviations here are about 1 / √(the axis each weight is multiplied along): 1/4 for the 16 values of a patch, 1/8 for the 64 hidden values, 0.18 for the 32 channels. That’s what Chapter 10’s train.py did by hand.

A model’s name is a type. Functions take a model as a parameter and read its fields with ., so forward loses five parameters. Replace forward, loss and predict with:

## The class scores of each image.
fn forward(m: Spotter, img: Images u8) -> [b C] Logit over C:
    x = embed(img |> to Unit |> to Signed, m.w_e)
    for l in 0..L:
        x = x |> mixer.mix_tokens(m.tw1[l], m.tw2[l]) |> mixer.mix_channels(m.cw1[l], m.cw2[l])
    return x |> mixer.norm |> head(m.w_h)

## The mean cross-entropy of the scores against the labels.
fn loss(m: Spotter, img: Images u8, y: [b] Index[C]) -> [] NonNeg:
    return cross_entropy(forward(m, img), y, over=C)

## The likeliest class of each image.
fn predict(m: Spotter, img: Images u8) -> [b] Index[C]:
    return argmax(forward(m, img), over=C)

m.tw1 is a [L N Dt], as declared, and m.tw1[l] is block l’s [N Dt], as before. Delete predict_saved, and the import python torch it needed: Section 3 saves and loads whole models, by name, with every shape checked, which beats w[3] as [L D Dc].

A field the model doesn’t have is an error, with a suggestion:

dim D = 4

model M:
    param w: [D D] = zeros

fn f(m: M, x: [b D]) -> [b D] = x @ m.W

$ shapescript check typo.ss
error[E0901]: `M` has no field `W`
 --> typo.ss:6:39
  |
3 | model M:
  |       - `M` is declared here
...
6 | fn f(m: M, x: [b D]) -> [b D] = x @ m.W
  |                                       ^
  |
  = help: a field with a similar name exists: `w`

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

Spotter.init(seed=0) makes a spotter, from ShapeScript or from Python. Each field draws its random numbers from a generator of its own, seeded from the model’s seed and the field’s name, so the same seed always gives the same spotter, and adding a field doesn’t change the others. From Python, a model is an ordinary torch.nn.Module with every field registered, so state_dict(), .parameters() and the rest of PyTorch work with it:

from spotter import Spotter, predict
from shapes import draw

m = Spotter.init(seed=0)
print(sum(p.numel() for p in m.parameters()))
img, y = draw(4)
print(predict(m, img))

$ shapescript run peek.py
12928
tensor([2, 2, 0, 2])

12,928 numbers, and untrained, it calls almost everything a plus. The rest of the chapter trains it.

Buffers

A model can also hold tensors that are saved with it but never trained: buffers, declared with buffer instead of param. Running statistics are the usual example. The spotter doesn’t need one, but a model that subtracts the average image from its inputs would keep that average as a buffer:

dim S = 16

model Centered:
    ## Subtracted from every image. Saved with the model, never trained.
    buffer center: [1 S S] = zeros

fn centered(m: Centered, img: [b 1 S S]) -> [b 1 S S] = img - m.center

Changing a model with mut

Functions so far have only read the spotter. To teach it a new set of shapes, you might keep the mixer blocks, which have learned to see strokes and corners, and start the head again from zero. That’s a change to the model, and a function can change a model only if its parameter is marked mut. Every call writes mut too, so the signature and the call both show what changes:

## Clears the head, to retrain it on new classes.
fn reset_head(mut m: Spotter):
    m.w_h = m.w_h * 0.0

Forgetting either mut is an error:

model Spotter:
    param w_h: [D C] = normal(std=0.18)

fn reset_head(m: Spotter):
    m.w_h = m.w_h * 0.0

fn reset_head2(mut m: Spotter):
    m.w_h = m.w_h * 0.0

fn retrain(m: Spotter):
    reset_head2(m)

$ shapescript check no_mut.ss
error[E1102]: `m` can't be changed in `reset_head`
 --> no_mut.ss:5:5
  |
4 | fn reset_head(m: Spotter):
  |               - `m` is declared here
5 |     m.w_h = m.w_h * 0.0
  |     ^ `m` isn't a `mut` parameter
  |
  = help: to change it, declare it `mut m` and pass it as `mut` at every call

error[E1103]: `reset_head2` changes `m`, so its argument needs `mut`
  --> no_mut.ss:11:17
   |
 7 | fn reset_head2(mut m: Spotter):
   |                -------------- `m` is declared `mut` here
...
11 |     reset_head2(m)
   |                 ^
   |
   = help: write `mut m`

For more information about E1102 and E1103, run `shapescript explain` with the code.

Only models and optimizer state can be mut. Tensors and numbers are values: a function that computes a new one returns it.

Gradients

value_and_grad(f, m, args...) calls f(m, args...) and returns two things: its value, and its gradient with respect to m’s params:

## The loss and its gradient.
fn grads(m: Spotter, img: Images u8, y: [b] Index[C]) -> Grad[Spotter]:
    l, g = value_and_grad(loss, m, img, y)
    return g

g has type Grad[Spotter]: a record with a field for each param, g.tw1: [L N Dt] and so on (and none for buffers, which aren’t trained). Gradients of one model can be added, subtracted and scaled, and norm(g) is the norm of all their values together.

The checker reads f to see which params it uses. A loss that doesn’t use the model at all would always have a zero gradient, so it’s an error; a param it never reads gets a warning, which is often a forgotten layer. Say a forward mixes the channels in each block, but forgets to mix the patches:

model Spotter:
    param w_e: [P*P D]  = normal(std=0.25)
    param tw1: [L N Dt] = normal(std=0.25)
    param tw2: [L Dt N] = normal(std=0.125)
    param cw1: [L D Dc] = normal(std=0.18)
    param cw2: [L Dc D] = normal(std=0.125)
    param w_h: [D C]    = normal(std=0.18)

fn forward(m: Spotter, img: Images u8) -> [b C] Logit over C:
    x = embed(img |> to Unit |> to Signed, m.w_e)
    for l in 0..L:
        x = x |> mixer.mix_channels(m.cw1[l], m.cw2[l])
    return x |> mixer.norm |> pool |> matmul(m.w_h)

fn loss(m: Spotter, img: Images u8, y: [b] Index[C]) -> [] NonNeg:
    return cross_entropy(forward(m, img), y, over=C)

fn grads(m: Spotter, img: Images u8, y: [b] Index[C]) -> Grad[Spotter]:
    l, g = value_and_grad(loss, m, img, y)
    return g

$ shapescript check forgot.ss
warning[W1101]: `loss` never reads the param `tw1`, so its gradient is always zero
  --> forgot.ss:19:27
   |
 3 |     param tw1: [L N Dt] = normal(std=0.25)
   |           --- `tw1` is declared here
...
19 |     l, g = value_and_grad(loss, m, img, y)
   |                           ^^^^
   |
   = help: remove the param, or make it a `buffer` if it isn't trained

warning[W1101]: `loss` never reads the param `tw2`, so its gradient is always zero
  --> forgot.ss:19:27
   |
 4 |     param tw2: [L Dt N] = normal(std=0.125)
   |           --- `tw2` is declared here
...
19 |     l, g = value_and_grad(loss, m, img, y)
   |                           ^^^^
   |
   = help: remove the param, or make it a `buffer` if it isn't trained

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

The program would run, and train, and do worse than it should, with nothing in the loss curve to say why. Here the checker says which weights are dead.

Optimizers

sgd(m, lr=...) and adamw(m, lr=...) make optimizer state for a model. Its type comes from the model’s: adamw(m, lr=3e-3) with m: Spotter is an AdamW[Spotter], which covers every param of Spotter, so none can be left out, as one could be from the list Chapter 10 gave torch.optim.Adam. update(mut m, mut opt, g) takes one step:

## One step of training on a batch, giving the loss before it.
fn step(mut m: Spotter, mut opt: AdamW[Spotter], img: Images u8, y: [b] Index[C]) -> [] NonNeg:
    l, g = value_and_grad(loss, m, img, y)
    update(mut m, mut opt, g |> clip_norm(1.0))
    return l

## `steps` steps on the same batch.
fn train(mut m: Spotter, img: Images u8, y: [b] Index[C], steps: int):
    opt = adamw(m, lr=warmup(20, cosine(3e-3, steps)))
    for i in 0..steps:
        l = step(mut m, mut opt, img, y)

• update changes both the model and the optimizer state, so both are mut.
• clip_norm(1.0) is a gradient transform: a function from Grad[Spotter] to Grad[Spotter], used with a pipe. It scales the gradient down when its norm is above 1, which keeps one bad batch from throwing the weights far. Your own transforms are ordinary functions, such as fn halve(g: Grad[Spotter]) -> Grad[Spotter] = g * 0.5.
• The learning rate can be a number or a schedule: warmup(20, then) rises from 0 over 20 steps, then follows then; cosine(peak, steps) decays from peak to 0. The optimizer state counts its own updates, so there’s no scheduler to step.

Optimizer state is matched to its model. Passing an AdamW[Other] to update a Spotter is an error, as is an SGD[Spotter] where an AdamW[Spotter] is expected.

train only ever sees one batch. Real training needs a dataset, shuffled into batches, which is the next section.

The spotter so far

devices cpu

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 spotter's weights.
model Spotter:
    param w_e: [P*P D]  = normal(std=0.25)
    param tw1: [L N Dt] = normal(std=0.25)
    param tw2: [L Dt N] = normal(std=0.125)
    param cw1: [L D Dc] = normal(std=0.18)
    param cw2: [L Dc D] = normal(std=0.125)
    param w_h: [D C]    = normal(std=0.18)

## 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(m: Spotter, img: Images u8) -> [b C] Logit over C:
    x = embed(img |> to Unit |> to Signed, m.w_e)
    for l in 0..L:
        x = x |> mixer.mix_tokens(m.tw1[l], m.tw2[l]) |> mixer.mix_channels(m.cw1[l], m.cw2[l])
    return x |> mixer.norm |> head(m.w_h)

## The mean cross-entropy of the scores against the labels.
fn loss(m: Spotter, img: Images u8, y: [b] Index[C]) -> [] NonNeg:
    return cross_entropy(forward(m, img), y, over=C)

## The likeliest class of each image.
fn predict(m: Spotter, img: Images u8) -> [b] Index[C]:
    return argmax(forward(m, img), 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)))

## Clears the head, to retrain it on new classes.
fn reset_head(mut m: Spotter):
    m.w_h = m.w_h * 0.0

## The loss and its gradient.
fn grads(m: Spotter, img: Images u8, y: [b] Index[C]) -> Grad[Spotter]:
    l, g = value_and_grad(loss, m, img, y)
    return g

## One step of training on a batch, giving the loss before it.
fn step(mut m: Spotter, mut opt: AdamW[Spotter], img: Images u8, y: [b] Index[C]) -> [] NonNeg:
    l, g = value_and_grad(loss, m, img, y)
    update(mut m, mut opt, g |> clip_norm(1.0))
    return l

## `steps` steps on the same batch.
fn train(mut m: Spotter, img: Images u8, y: [b] Index[C], steps: int):
    opt = adamw(m, lr=warmup(20, cosine(3e-3, steps)))
    for i in 0..steps:
        l = step(mut m, mut opt, img, y)

## 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. Write sgd_head(mut m: Spotter, img: Images u8, y: [b] Index[C], lr: float), which updates only the head by hand, without an optimizer: w_h minus lr times its gradient, leaving the mixer blocks as they are. That’s how you’d retrain the head after reset_head.

2. Find the bugs:

   dim D = 32
   dim C = 4
   
   model Head:
       param w: [D C] = normal(std=0.18)
   
   fn loss(m: Head, x: [b D], y: [b] Index[C]) -> [] = cross_entropy(x @ m.w, y, over=C)
   
   fn step(m: Head, opt: SGD[Head], x: [b D], y: [b] Index[C]):
       l, g = value_and_grad(loss, m, x, y)
       update(m, opt, g)
   

3. Write a gradient transform scale_down(g: Grad[Spotter], k: float) -> Grad[Spotter], and use it in step’s pipe after clip_norm, to take steps half as long.

Data, Modes and Logs

The spotter has weights and can take a step. To train it properly it needs three more things: some randomness while training, a dataset to draw batches from, and a way to report how it’s going.

Modes and dropout

A spotter that sees the same 2,000 pictures again and again can learn those pictures by heart instead of learning shapes. Dropout is the usual guard: while training, it zeroes a random tenth of the values, so the head can’t rely on any one of them. When the spotter is used, it must be off.

In PyTorch, model.train() and model.eval() flip a flag that dropout reads, and forgetting eval() before evaluating is a classic bug. In ShapeScript there’s no flag. A function that does something random takes a Mode argument, which is Train(seed=s) or Eval. Change forward to take one, and put dropout between the pooling and the head’s matrix product; head can go, since forward now does its two steps with dropout between them:

## The class scores of each image.
fn forward(m: Spotter, img: Images u8, mode: Mode) -> [b C] Logit over C:
    x = embed(img |> to Unit |> to Signed, m.w_e)
    for l in 0..L:
        x = x |> mixer.mix_tokens(m.tw1[l], m.tw2[l]) |> mixer.mix_channels(m.cw1[l], m.cw2[l])
    return x |> mixer.norm |> pool |> dropout(0.1, mode) |> matmul(m.w_h)

## The likeliest class of each image.
fn predict(m: Spotter, img: Images u8) -> [b] Index[C]:
    return argmax(forward(m, img, Eval), over=C)

predict passes Eval, so whoever calls it can’t get it wrong. With Eval, dropout returns its input; with Train(seed=s), it zeroes each value with probability 0.1 and divides the others by 0.9.

A mode isn’t a flag. A bool named training would be the PyTorch habit, and the checker refuses it:

dim D = 32
dim C = 4

fn scores(x: [b D], w: [D C], training: bool) -> [b C]:
    return x |> dropout(0.1, training) |> matmul(w)

$ shapescript check flag.ss
error[E1202]: expected a `Mode`, found `bool`
 --> flag.ss:5:30
  |
5 |     return x |> dropout(0.1, training) |> matmul(w)
  |                              ^^^^^^^^ `dropout`'s parameter `mode` is a `Mode`
  |
  = help: pass `Train(seed=…)` to train or `Eval` to evaluate
  = note: in stage 1 of this pipeline, whose input is `[b D]`

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

The difference is the seed. A mode carries the seed that random operations draw from, so training is reproducible with no global random state. Each call that passes a mode on folds in its own name, so the two dropouts of a model with two draw different masks from one seed. And randomness shows in the signatures: dropout needs a mode, so forward takes one, and so do loss and everything else that calls it, with each call passing one on. A function without a Mode parameter, like predict, either passes Eval or isn’t random.

Data records

Until now, a batch was two arguments, img and y, that had to agree on b. A data declaration says what one example holds, so the two travel together:

## One labelled picture.
data Shape:
    img:   [1 S S] u8
    label: [] Index[C]

[b] Shape is a batch of b examples. Every field gets the batch’s axis first: batch.img is [b 1 S S] u8, an Images u8, and batch.label is [b] Index[C]. So loss can take one batch instead of two tensors, and pass the mode on:

## The mean cross-entropy of the scores against the labels.
fn loss(m: Spotter, batch: [b] Shape, mode: Mode) -> [] NonNeg:
    return cross_entropy(forward(m, batch.img, mode), batch.label, over=C)

## The fraction of a batch predicted right.
fn accuracy(m: Spotter, batch: [b] Shape) -> []:
    return mean((predict(m, batch.img) == batch.label) to f32)

and step changes to match. Delete grads, which was there to show a gradient, and train, which only ever saw one batch:

## One step of training on a batch, giving the loss before it.
fn step(mut m: Spotter, mut opt: AdamW[Spotter], batch: [b] Shape, mode: Mode) -> [] NonNeg:
    l, g = value_and_grad(loss, m, batch, mode)
    update(mut m, mut opt, g |> clip_norm(1.0))
    return l

A misspelled field is an error, and so is a field whose shape uses a generic size: every example of a record has the same shapes, so they can be stacked into batches. (Examples of different lengths, such as sentences, are padded to a dim, with a Mask field saying which positions are real. Padding them for you is coming.)

Datasets and batches

Dataset[Shape] is a dataset of Shapes. It comes from Python: any map-style dataset, an object with __len__ and __getitem__ such as a torch.utils.data.Dataset, whose items have the record’s fields, as a dict or a tuple. Section 3 wraps draw in one. Each item is checked against the record as it’s loaded, so a picture of the wrong size, or a label of 4, fails at once, naming the field.

Datasets are shuffled and batched with two pipe stages, and a for loop runs over the batches. Here’s how well the spotter does on each batch of a dataset:

## Prints the accuracy of each batch, in a shuffled order.
fn show_accuracy(m: Spotter, data: Dataset[Shape], epoch: int):
    for batch in data |> shuffle(seed=epoch) |> batches(64):
        print("accuracy", accuracy(m, batch))

shuffle(seed=epoch) gives an order decided by the seed alone. The last batch holds what’s left over, so the batch size isn’t known statically: inside the loop, batch is a [b] Shape whose b is an unknown size. (When it can, the loader pads the last batch to the full size and masks the padding out, so every step has the same shapes; means and losses still count only the real examples.)

Batches are built by worker processes and copied to the declared device for you.

Logs and sinks

A logs declaration names the numbers a program reports, and how each combines when several are reported together, such as over an epoch:

## What each training step reports.
logs StepLogs:
    loss:      [] mean
    accuracy:  [] mean
    grad_norm: [] max

Over an epoch, a row will have the mean loss, the mean accuracy, and the largest gradient norm, the one to watch for a step that went wrong. StepLogs(loss=l, accuracy=a, grad_norm=n) makes a record, and every field must be given, by name. A misspelled metric is a compile error, not a missing column you notice a week later:

fn report(l: [] NonNeg) -> StepLogs = StepLogs(loss=l, accuracy=0.0, grad_nrom=0.0)

Records are written to a sink, a LogSink, which comes from Python: any object with a write(dict) method. The shapescript.sinks module has sinks for CSV files and TensorBoard that combine records by each field’s rule: shapescript.sinks.csv("log.csv", every=None) writes one row each time the program calls sink.flush().

## Reports one step.
fn log_step(sink: LogSink, l: [] NonNeg, a: [], n: [] NonNeg):
    sink.write(StepLogs(loss=l, accuracy=a, grad_norm=n))

print(a, b, ...), as in show_accuracy, prints any values, tensors included, copying them to the CPU as needed.

The spotter so far

devices cpu

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 spotter's weights.
model Spotter:
    param w_e: [P*P D]  = normal(std=0.25)
    param tw1: [L N Dt] = normal(std=0.25)
    param tw2: [L Dt N] = normal(std=0.125)
    param cw1: [L D Dc] = normal(std=0.18)
    param cw2: [L Dc D] = normal(std=0.125)
    param w_h: [D C]    = normal(std=0.18)

## One labelled picture.
data Shape:
    img:   [1 S S] u8
    label: [] Index[C]

## What each training step reports.
logs StepLogs:
    loss:      [] mean
    accuracy:  [] mean
    grad_norm: [] max

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

## 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(m: Spotter, img: Images u8, mode: Mode) -> [b C] Logit over C:
    x = embed(img |> to Unit |> to Signed, m.w_e)
    for l in 0..L:
        x = x |> mixer.mix_tokens(m.tw1[l], m.tw2[l]) |> mixer.mix_channels(m.cw1[l], m.cw2[l])
    return x |> mixer.norm |> pool |> dropout(0.1, mode) |> matmul(m.w_h)

## The mean cross-entropy of the scores against the labels.
fn loss(m: Spotter, batch: [b] Shape, mode: Mode) -> [] NonNeg:
    return cross_entropy(forward(m, batch.img, mode), batch.label, over=C)

## The likeliest class of each image.
fn predict(m: Spotter, img: Images u8) -> [b] Index[C]:
    return argmax(forward(m, img, Eval), over=C)

## The fraction of a batch predicted right.
fn accuracy(m: Spotter, batch: [b] Shape) -> []:
    return mean((predict(m, batch.img) == batch.label) to f32)

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

## Clears the head, to retrain it on new classes.
fn reset_head(mut m: Spotter):
    m.w_h = m.w_h * 0.0

## One step of training on a batch, giving the loss before it.
fn step(mut m: Spotter, mut opt: AdamW[Spotter], batch: [b] Shape, mode: Mode) -> [] NonNeg:
    l, g = value_and_grad(loss, m, batch, mode)
    update(mut m, mut opt, g |> clip_norm(1.0))
    return l

## Prints the accuracy of each batch, in a shuffled order.
fn show_accuracy(m: Spotter, data: Dataset[Shape], epoch: int):
    for batch in data |> shuffle(seed=epoch) |> batches(64):
        print("accuracy", accuracy(m, batch))

## Reports one step.
fn log_step(sink: LogSink, l: [] NonNeg, a: [], n: [] NonNeg):
    sink.write(StepLogs(loss=l, accuracy=a, grad_norm=n))

## 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. Pictures in colour would need a different record. Write data ColorShape with an RGB image, [3 S S] u8, and the label, and give the type of batch.img for batch: [b] ColorShape.
2. Write count_right(m: Spotter, batch: [b] Shape) -> [], the number of pictures of a batch the spotter gets right.
3. Add a field lr to StepLogs, combined by its last value with last, and write log_lr(sink: LogSink, l: [] NonNeg, lr: float) that reports it with the loss. Which rule would you give a learning rate, and why not mean?

The Training Loop and Checkpoints

This section puts the pieces together: a training loop in spotter.ss, checkpoints to save and resume it, and a short Python script that only provides the pictures. Then it makes the loop work for more than one model, and swaps it for the standard library’s.

The program

First, one step of training that also reports. It replaces Section 2’s step:

## One step of training.
fn train_step(mut m: Spotter, mut opt: AdamW[Spotter], batch: [b] Shape, mode: Mode) -> StepLogs:
    l, g = value_and_grad(loss, m, batch, mode)
    n = norm(g)
    update(mut m, mut opt, g |> clip_norm(1.0))
    return StepLogs(loss=l, accuracy=accuracy(m, batch), grad_norm=n)

Delete show_accuracy and log_step from Section 2: train_step and the loop do their jobs. Then the loop, with checkpoints:

## Trains a new spotter for `epochs` epochs.
fn fit(mut m: Spotter, train: Dataset[Shape], sink: LogSink, epochs: int, mode: Mode):
    opt = adamw(m, lr=3e-3)
    run(mut m, mut opt, train, sink, 0, epochs, mode)

## Carries on from the checkpoint of the last epoch that finished.
fn resume(train: Dataset[Shape], sink: LogSink, epochs: int, mode: Mode) -> Spotter:
    m, opt, done = load("ckpt.pt", Spotter, AdamW[Spotter], int)
    run(mut m, mut opt, train, sink, done, epochs, mode)
    return m

## Epochs `first` to `epochs`, with a checkpoint after each.
fn run(mut m: Spotter, mut opt: AdamW[Spotter], train: Dataset[Shape], sink: LogSink,
                first: int, epochs: int, mode: Mode):
    for epoch in first..epochs:
        for batch in train |> shuffle(seed=epoch) |> batches(64):
            sink.write(train_step(mut m, mut opt, batch, mode))
        sink.flush()
        save("ckpt.pt", m, opt, epoch + 1)
        print("epoch", epoch, "done")

Read run from the top: each epoch shuffles with its own seed, writes a log record per step, flushes the sink (one CSV row per epoch), and saves a checkpoint. Nothing happens that the code doesn’t say. That’s the whole of Chapter 10’s Python loop, the optimizer included, now checked: train_step can’t be given an optimizer for another model, a batch of the wrong record, or a misspelled metric.

Checkpoints

save(path, values...) writes models, optimizer state (moments, learning rate, step count) and host numbers, such as the number of epochs done. load(path, Types...) takes the type of each value in the file and checks the file against them before it makes anything: a Spotter saved before you changed D is an error naming each field whose shape differs, not a w[0] as [P*P D] that fails on the first call. Loaded tensors go on the program’s device, never the one they were saved from.

Only models, optimizer state and numbers can be saved. A tensor that’s part of the training state belongs in a model, as a buffer if it isn’t trained:

dim S = 16

fn keep(img: [b 1 S S]):
    save("img.pt", img)

Running it

The Python side is now just the pictures. Shapes wraps Chapter 7’s draw as a dataset, with each item a dict of the record’s fields:

import shapescript
from shapes import draw
from spotter import Spotter, fit, resume, predict


class Shapes:
    """`n` labelled pictures from `draw`, for ShapeScript's loader."""

    def __init__(self, n, seed):
        self.img, self.label = draw(n, seed)

    def __len__(self):
        return len(self.label)

    def __getitem__(self, i):
        return {"img": self.img[i], "label": int(self.label[i])}


shapescript.loader_workers(0)  # the dataset is defined in this script
train = Shapes(2000, seed=0)
sink = shapescript.sinks.csv("log.csv", every=None)
model = Spotter.init(seed=0)
fit(model, train=train, sink=sink, epochs=4, mode=shapescript.Train(seed=0))
model = resume(train=train, sink=sink, epochs=8, mode=shapescript.Train(seed=0))
sink.close()

test = Shapes(500, seed=1)
right = (predict(model, test.img) == test.label).float().mean().item()
print(f"test accuracy: {right:.1%}")

$ shapescript run train.py
epoch 0 done
epoch 1 done
epoch 2 done
epoch 3 done
epoch 4 done
epoch 5 done
epoch 6 done
epoch 7 done
test accuracy: 97.2%
$ cat log.csv
loss,accuracy,grad_norm
1.0460098292678595,0.59521484375,1.4748300313949585
0.501040487550199,0.8466796875,2.4373481273651123
0.26240775594487786,0.91845703125,2.432314395904541
0.1644725656369701,0.962890625,2.5464417934417725
0.08703674498246983,0.984375,2.1427037715911865
0.040361865219892934,0.99560546875,1.6711437702178955
0.028169181416160427,0.99853515625,1.5459754467010498
0.010340657085180283,1.0,1.0254361629486084

fit trained four epochs and saved after each; resume loaded the fourth checkpoint and carried on to eight. The result is exactly that of training eight epochs without stopping, dropout included: each dropout mask is drawn from the mode, the epoch and the batch, not from how many steps this run has taken, so a run that starts at epoch 4 draws the masks of epoch 4. fit(model, ...) changed model in place: from Python, a mut argument is changed, as a PyTorch module would be. loader_workers(0) loads batches in this process, because a dataset class defined in the script you run can’t be sent to worker processes.

Read the log’s last row against the test accuracy. In the last epoch, the spotter got every training picture right, but nearly 3% of new ones wrong. It has started to learn its 2,000 pictures by heart, which is what dropout is there to slow down; more pictures, or the nudged ones of Chapter 10, would help more.

One loop for every model

train_step only trains a Spotter. Is the spotter worth its mixer blocks? A fair test trains a much simpler model on the same batches. Two features make one loop work for both.

Overloading by model. Several functions may share a name if each takes a different model first. So each model brings its own forward and loss, and a call picks the one for its model. In a new file, compare.ss, next to spotter.ss:

import spotter

## A baseline: one matrix from pixels to scores.
model Linear:
    param w: [spotter.S*spotter.S spotter.C] = zeros

fn loss(m: Linear, batch: [b] spotter.Shape, mode: Mode) -> [] NonNeg:
    x = batch.img |> to Unit |> rearrange(b c h w -> b (c h w))
    return cross_entropy(x @ m.w, batch.label, over=spotter.C)

loss now has two versions, spotter.loss for a Spotter and this one for a Linear.

Generic functions. A parameter type that’s a single capital letter, such as M or E, that nothing declares, is a type variable: m: M takes any model, and [b] E a batch of any data record. A function with one is checked again at each call, for the types it’s given, so it needs only what it uses. There are no traits or interfaces to declare:

## One SGD step on any model that has a `loss`.
fn step(mut m: M, mut opt: SGD[M], batch: [b] E, mode: Mode) -> []:
    l, g = value_and_grad(loss, m, batch, mode)
    update(mut m, mut opt, g)
    return l

## The losses of both models, trained side by side for an epoch.
fn race(mut a: spotter.Spotter, mut z: Linear, data: Dataset[spotter.Shape], mode: Mode):
    oa = sgd(a, lr=0.1)
    oz = sgd(z, lr=0.1)
    for batch in data |> shuffle(seed=0) |> batches(64):
        print(step(mut a, mut oa, batch, mode), step(mut z, mut oz, batch, mode))

At step(mut a, ...), M is spotter.Spotter and loss(m, ...) is the spotter’s; at step(mut z, ...), it’s Linear’s. A model without the function a generic one needs is an error at the call, naming what’s missing and where it’s used.

The standard library’s fit

The standard library has a generic fit of its own, written in ShapeScript like step above, with the batch size and learning rate as parameters. All it needs from a model is a loss(m, batch, mode), which the spotter has. Import it at the top of spotter.ss:

import shapescript.std.fit as std

and call it:

## Trains with the standard library's loop instead.
fn fit_std(mut m: Spotter, train: Dataset[Shape], sink: LogSink, epochs: int, mode: Mode):
    std.fit(mut m, train, sink, epochs, 64, 3e-3, mode)

It’s short enough to read on one screen, and meant to be copied when you need something different: shapescript.std.path("fit") in Python gives the file. Read it, and you’ll find a train_step and a run, the names of the spotter’s own. They don’t clash, because each file is a namespace. Inside fit.ss, run means fit.ss’s run, whatever the model’s file declares. A generic function looks next to the model only for names its own file doesn’t declare, which is how std.fit finds the spotter’s loss. So you don’t need to know a library’s helpers to name your functions, and a library can add helpers without changing what your code means.

The spotter so far

The whole of spotter.ss, as the project leaves it:

devices cpu

import mixer
import shapescript.std.fit as std

## 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 spotter's weights.
model Spotter:
    param w_e: [P*P D]  = normal(std=0.25)
    param tw1: [L N Dt] = normal(std=0.25)
    param tw2: [L Dt N] = normal(std=0.125)
    param cw1: [L D Dc] = normal(std=0.18)
    param cw2: [L Dc D] = normal(std=0.125)
    param w_h: [D C]    = normal(std=0.18)

## One labelled picture.
data Shape:
    img:   [1 S S] u8
    label: [] Index[C]

## What each training step reports.
logs StepLogs:
    loss:      [] mean
    accuracy:  [] mean
    grad_norm: [] max

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

## 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(m: Spotter, img: Images u8, mode: Mode) -> [b C] Logit over C:
    x = embed(img |> to Unit |> to Signed, m.w_e)
    for l in 0..L:
        x = x |> mixer.mix_tokens(m.tw1[l], m.tw2[l]) |> mixer.mix_channels(m.cw1[l], m.cw2[l])
    return x |> mixer.norm |> pool |> dropout(0.1, mode) |> matmul(m.w_h)

## The mean cross-entropy of the scores against the labels.
fn loss(m: Spotter, batch: [b] Shape, mode: Mode) -> [] NonNeg:
    return cross_entropy(forward(m, batch.img, mode), batch.label, over=C)

## The likeliest class of each image.
fn predict(m: Spotter, img: Images u8) -> [b] Index[C]:
    return argmax(forward(m, img, Eval), over=C)

## The fraction of a batch predicted right.
fn accuracy(m: Spotter, batch: [b] Shape) -> []:
    return mean((predict(m, batch.img) == batch.label) to f32)

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

## Clears the head, to retrain it on new classes.
fn reset_head(mut m: Spotter):
    m.w_h = m.w_h * 0.0

## One step of training.
fn train_step(mut m: Spotter, mut opt: AdamW[Spotter], batch: [b] Shape, mode: Mode) -> StepLogs:
    l, g = value_and_grad(loss, m, batch, mode)
    n = norm(g)
    update(mut m, mut opt, g |> clip_norm(1.0))
    return StepLogs(loss=l, accuracy=accuracy(m, batch), grad_norm=n)

## Trains a new spotter for `epochs` epochs.
fn fit(mut m: Spotter, train: Dataset[Shape], sink: LogSink, epochs: int, mode: Mode):
    opt = adamw(m, lr=3e-3)
    run(mut m, mut opt, train, sink, 0, epochs, mode)

## Carries on from the checkpoint of the last epoch that finished.
fn resume(train: Dataset[Shape], sink: LogSink, epochs: int, mode: Mode) -> Spotter:
    m, opt, done = load("ckpt.pt", Spotter, AdamW[Spotter], int)
    run(mut m, mut opt, train, sink, done, epochs, mode)
    return m

## Epochs `first` to `epochs`, with a checkpoint after each.
fn run(mut m: Spotter, mut opt: AdamW[Spotter], train: Dataset[Shape], sink: LogSink,
                first: int, epochs: int, mode: Mode):
    for epoch in first..epochs:
        for batch in train |> shuffle(seed=epoch) |> batches(64):
            sink.write(train_step(mut m, mut opt, batch, mode))
        sink.flush()
        save("ckpt.pt", m, opt, epoch + 1)
        print("epoch", epoch, "done")

## Trains with the standard library's loop instead.
fn fit_std(mut m: Spotter, train: Dataset[Shape], sink: LogSink, epochs: int, mode: Mode):
    std.fit(mut m, train, sink, epochs, 64, 3e-3, mode)

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

From six loose tensors and a Python loop to a model that trains itself: that’s the spotter done. Chapters 13 to 16 are about the tools around a program like it, and Chapter 17 builds a second project, a small transformer, with everything at once.

Exercises

1. Make run print "saved" after every checkpoint. Where does the line go, and why must it be indented like save?
2. Write retrain_head(mut m: Spotter, train: Dataset[Shape], sink: LogSink, epochs: int, mode: Mode), which clears the head with reset_head and trains again with fit. Which muts does it need?
3. Train compare.ss’s Linear baseline with std.fit, in a function fit_linear in compare.ss. What does std.fit need from Linear that it already has?

The Editor

ShapeScript was designed for the editor as much as for the command line. Its language server, shapescript-lsp, gives every LSP editor errors as you type, hovers, navigation and more, and the VS Code extension adds pictures: shape diagrams, activation curves and model architectures, drawn as SVG. Chapter 1 shows how to set up VS Code: installing the extension, which is coming to the Marketplace, and selecting your project’s Python. If something there doesn’t work, see Appendix D.

This chapter tours what the editor shows. Open examples/attention.ss and try each one as you read.

Errors as you type

Every diagnostic in this book appears in the editor as you type, underlined where the checker points. Hovering the underline shows the message, and its code links to the reference. Fix it and the underline goes.

Hover: shapes as diagrams

Hover any expression, any pipe stage or any name, and you see its type, drawn:

• 1-, 2- and 3-axis tensors are a bar, a rectangle or a box, with each side labeled.
• More axes are a stack of copies of the last two, labeled with the rest (× H × b).
• Grouped axes, such as (H Dh), are drawn as segmented edges.
• Masks show their structure. A causal mask is drawn as a triangle, in VS Code as SVG and elsewhere as text:

mask: [n m] Mask  ·  lower-triangular  ·  m ≡ n
         m
    ◄─────────►
  ▲ ┌─────────┐
  │ │█        │
  n │███      │
  │ │█████    │
  ▼ │███████  │
    └─────────┘

Hovering a dim shows its value and doc comment, then where it’s used: the values whose types mention it, grouped by function. Hovering an axis name shows where it was introduced. Hovering a built-in that moves values, such as rearrange or split, explains it, with an example, and then shows this call’s tensors before and the result after, side by side. Hovering |>, or -> in a pattern, does the same for that step.

Hover: what a function does

Hover a function’s name, or a call such as argmax(z, over=C), and the editor runs it on a small example and draws what happens. The example comes from the types, with the dims made small, so your functions get pictures as the built-ins do. When the result’s type says something about its values, such as Prob over m or Index[C], the values going in and coming out are drawn: argmax puts a 1 where each row is largest, and softmax makes rows that add up to 1. Otherwise, the picture shows what depends on what. One value going in is changed, and the values coming out that change are outlined. For the causal_attention of the Introduction, VS Code draws it as an image; other editors show it as text:

x: [b n D]        [b n D]
  D                 D
  · · · · · ·       · · · · · ·
  · · · · · ·       · · · · · ·
n · · ● · · · ──→ n ● ● ● ● ● ●
  · · · · · ·       ● ● ● ● ● ●
  · · · · · ·       ● ● ● ● ● ●
at b = 0          at b = 0

Changing position 2 of x changes the output at positions 2 and after, but not before: the causal mask, seen from outside. The weights w_qkv and w_o aren’t drawn. The editor runs the function in your Python interpreter, which needs PyTorch, only in a trusted workspace, and only if the function doesn’t print, save, load, write logs, call Python or change a mut parameter.

Shapes at the end of each line

In a pipeline written one stage per line, each line ends with its shape, as an inlay hint, updated as you type. Hover the hint for the diagram.

Activation curves

Hover an activation, such as gelu, relu, silu or sigmoid, and you see its formula, its curve with the derivative dashed, and its range. It works for your own functions too, wherever they’re declared, called or piped into: an elementwise function, or any function of one tensor and some numbers that only does elementwise maths to the tensor:

## Squared ReLU, scaled.
fn relu2(x: [b D], k: float = 1.0) -> [b D] = square(relu(x)) * k

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

fn layer(h: [b 64]) -> [b 64]:
    return h |> relu2(0.5) |> swish

Hovering relu2 in layer plots it with k = 0.5, the number written at the call. VS Code draws the curve as SVG; other editors draw it in braille.

Math in comments

Comments can hold LaTeX: $...$ inline, $$...$$ on lines of its own. Hovering shows it typeset in VS Code, and as Unicode text elsewhere; math in a ## comment also appears in the hover of the name it documents. Typing \ offers the supported commands.

## Root-mean-square normalization:
## $$\RMSNorm(x) = \frac{x}{\sqrt{\epsilon + \frac{1}{D}\sum_{i=1}^{D} x_i^2}} \odot g$$
fn rms_norm(x: [b n D], g: [D], eps = 1e-6) -> [b n D]:
    # $\epsilon$ keeps the denominator away from 0.
    return x / sqrt(mean(x * x, over=D) + eps) * g

Math never causes an error. A command outside the supported subset is shown as written, with a hint:

## Normalizes with $\sofmax(x)$.
fn f(x: [b n]) -> [b n] = softmax(x, over=n)

A model’s architecture

Click a model’s name, where it’s declared or used as a type, and VS Code draws its architecture in an Architecture tab beside the code, which follows the cursor to other models (the CodeLens Show architecture above each model line opens it again after you close it). Hovering the name shows a one-line summary, such as Embedding → Transformer block × 6 → LayerNorm → Linear · tied. The architecture is read from the functions that take it first: its forward, then its loss. It finds the layers however the code writes them: linear layers, MLPs, attention, embeddings, norms, convolutions, dropout and losses, each box labeled with the type of its result and the params it holds, with residual connections drawn as arrows. A loop over layers is drawn once, marked × L. Hover Gpt in examples/gpt.ss to see a whole transformer, or the model of Chapter 17. Other editors get the same as a list:

forward  mlp.ss:26
  • x                                    [b D]
  ↓ MLP D → Hd → C · GELU · dropout 0.1  [b C]
        w1 [D Hd], w2 [Hd C]
  ↓ output                               [b C] Logit over C

In the tab, Every operation draws each operation instead of layers, and clicking a box selects the code it came from; the tab zooms and pans for big models. Other editors show the drawing in the hover.

Completion

The editor completes names as you type, everywhere a name can go, and puts first the ones whose types fit. Inside mix(, the values that are [b n D], which mix takes, come first, the first one already selected; then those that would also fit; then the rest, with what wouldn’t type-check last. Each item says what it is, as y: [b n D] or fn softmax(x: [... n] Logit over n, over: Axis[n]) -> [... n] Prob over n, and shows the same documentation as hovering it. Choosing a function writes the call, with its arguments to fill in: softmax(x, over=n), then Tab to go from one to the next.

Some places know more:

• After over=, the list is exactly the axes of the tensor the call works on: for a [b n D], b, n and D.
• After |>, it’s the functions that can take the value piped in, as stages: softmax(over=…), offering the value’s axes, or gelu alone. Functions of yours that take exactly that shape come before the built-ins.
• In rearrange(x, and the other pattern operations, it’s the input’s axes, all of them first, then one at a time; after ->, the names on the left.
• In a type, it’s the dims, the axes the signature already names, and after the shape, domains and dtypes. After as, the value’s own type comes first.
• After ., it’s what’s there: a batch’s fields with their types, an imported file’s declarations, init after a model’s name, and a Python module’s names with their signatures (your own modules are read, never run). After m., a model’s fields and the functions that take it, as calls (forward(m, …)).
• After import , the .ss files you can import; after import python , the Python modules.
• At the start of a line, the keywords that can go there: fn, model, data and the other declarations at the top level, written out with their blocks, and for, if, else and return in a function.

Matching is fuzzy: lnorm finds layer_norm. A name from another file of your project that isn’t imported yet is offered too, after the others, as metrics.accuracy: choosing it adds import metrics at the top.

While you type the arguments of a call, its signature shows above the cursor with the parameter you’re on highlighted, its type included. For overloaded functions, all of them are listed, the one your first argument picks first.

Navigation

• Go to definition works on axes, dims, functions, models, fields and imports, across files. On a built-in such as gelu, it opens the read-only declaration from the prelude. On torch.load, an import python module or an extern fn’s path, it opens the Python source.
• Find references and rename work on axes, dims, domains and functions. Renaming an axis changes it in the signature, the patterns and the expressions; a rename that would change what a name means is refused.
• Completion after m., for a model m, offers its fields and every function that takes the model first, even from files not imported yet. Choosing m.forward writes forward(m, …), and adds the import if one is needed.

Python editors know your .ss files

The extension writes a Python module beside each .ss file that checks, model.py beside model.ss, with typed declarations of its dims, models and functions and their docs. So in a Python file, Pylance and mypy resolve from model import forward, complete it, and show its ShapeScript signature on hover. (shapescript pymodules *.ss writes them from the command line.) A .py file you wrote yourself is never replaced.

Debugging

ShapeScript is debugged with Python’s debugger, debugpy, which VS Code and Jupyter use. Breakpoints in .ss files stop, stepping goes from ShapeScript into the Python it calls and back, and variables show their ShapeScript types and sizes:

x = [b=2 D=4] float32 tensor([[1., …

In VS Code, ShapeScript: Debug File in the Command Palette, or in the ▷ menu above the editor, debugs the current .ss file’s fn main(), or a Python program that imports .ss files, with the import hook installed. To keep a configuration, with arguments, add one of the extension’s shapescript launch configurations to .vscode/launch.json (Run › Add Configuration…). shapescript run runs a program without the debugger.

The debugger runs in the Python you selected for the project (Chapter 1), with the debugpy that comes with Microsoft’s Python extension. Without that extension, install debugpy into the environment, pip install debugpy, and name its Python in the setting shapescript.debug.python. The worker processes of several devices (Chapter 16) can’t be debugged yet.

To look at a tensor rather than read it, right-click it in the Variables view and choose View Tensor. It appears in the Tensor section of the Run and Debug view, under Variables and Watch, as a heatmap you can zoom from the whole tensor down to its numbers, arranged by its type: attention weights [b H n m] Prob over m as one heatmap per head, an image [b 3 h w] Unit in color. Beside it are checks of what the type promises, such as no NaN, probabilities that sum to 1, and zeros above the diagonal after a causal mask, each with Show to go to a value that breaks it. The program computes what the view draws on the GPU, so the tensor can be as big as the GPU holds.

Exercises

1. Open examples/attention.ss and hover m in softmax(over=m). Where does the editor say m comes from?

2. In curves.ss above, change relu2’s body to square(relu(x)) * k + mean(x). Does relu2 still get a curve? Why?

3. Fix the math so the hint goes away:

   ## The mean, $\fraq{1}{n}\sum_i x_i$.
   fn avg(x: [n]) -> [] = mean(x, over=n)
   

Notebooks

ShapeScript works in Jupyter and VS Code notebooks through a cell magic. Exploration is where shape bugs are born, so having the checker there helps most.

Setting up

A notebook runs its cells in a kernel, a Python process, and for ShapeScript cells it must be a Python with ShapeScript installed: your project’s environment. Chapter 1 sets it up (Notebooks): pip install ipykernel in the environment, then, in VS Code, Select Kernel at the top right of the notebook › Python Environments › .venv. In Jupyter in the browser, install jupyterlab into the environment and start jupyter lab from it.

To check which Python a notebook runs, run import sys; sys.executable in a cell: it should be inside your .venv. If %load_ext shapescript says No module named 'shapescript', it isn’t; see Appendix D.

Loading the magic

In a Python cell:

%load_ext shapescript

From then on, a cell that starts with %%shapescript is ShapeScript. A notebook declares its devices like a program, once:

%%shapescript
devices cpu
dim D = 768
fn mlp(x: [b D], w1: [D 4*D], w2: [4*D D]) -> [b D] = x @ w1 |> gelu |> matmul(w2)

Cells share what they declare

Unlike a file, a cell may also hold statements outside any function, and they run in order. Each cell is checked together with the cells that ran before it, so D and mlp from the cell above are known here:

%%shapescript
import python torch
x = torch.randn(8, 768) as [8 D]
y = mlp(x, torch.randn(768, 3072) as [D 4*D], torch.randn(3072, 768) as [4*D D])
y

The last line is a tensor, so the cell shows its shape diagram next to its value. Python variables of the notebook are visible too, as Python values, so a tensor made in a Python cell comes in with an as:

%%shapescript
z = mlp(batch as [4 D], torch.randn(768, 3072) as [D 4*D], torch.randn(3072, 768) as [4*D D])

(Here batch was made in a Python cell, batch = torch.randn(4, 768).) And ShapeScript values are Python variables afterwards, so a Python cell can plot y or save it.

A cell with errors runs nothing

A cell with a type error shows the error as its output, and runs nothing, so a half-run cell never leaves the notebook in a confusing state:

%%shapescript
w = torch.randn(768, 10) as [D 10]
bad = mlp(y @ w, w, w)

The output names [1], [2] and so on for the cells, in the order they ran, where a file would have a file name.

Running a cell again

Running a cell again replaces what it defined. If you change dim D = 768 to dim D = 512 and run that cell, later cells check against the new D the next time they run. Values made with the old D, such as y, can’t be used until the cell that made them runs again: their types no longer match the code, which is error E0701, rather than a silent mismatch.

Devices are the exception: a notebook’s devices can’t change once declared. Restart the kernel to change them.

Debugging in notebooks

Jupyter’s debugger, which VS Code and JupyterLab use, stops at breakpoints in %%shapescript cells, and steps into the Python they call, as it does for .ss files (Chapter 13).

Exercises

1. In a notebook where one cell declared dim D = 768, write a cell that declares fn norm(x: [b D]) -> [b D], a layer norm over D, and a second cell that calls it on a Python tensor t of shape (2, 768).
2. What happens if you run the second cell before the first?

Exporting without Python

A trained model often has to run where your Python code can’t: a server without your repository, a C++ application, a phone. PyTorch’s answer is torch.export, which records a function’s torch operations into a program that runs on its own. Its hard part is dynamic shapes: you must declare by hand which input sizes may change and how they relate, and you find out about mistakes when tracing fails.

ShapeScript’s types already say all of that. So shapescript export builds the dynamic shapes from the signature.

export fn

Mark a function export:

dim D = 768

## Splits each image into 16×16 patches.
export fn patchify(img: [b 3 16gh 16gw]) -> [b (gh gw) D]:
    return img |> rearrange(b c (gh p1) (gw p2) -> b (gh gw) (p1 p2 c))

and export it:

shapescript export patches.ss -o patchify.pt2

Exporting runs torch.export in a Python with PyTorch and einops (--python, or the active virtualenv’s), but the program it writes needs neither ShapeScript nor your code:

import torch
patchify = torch.export.load("patchify.pt2").module()   # no ShapeScript needed
print(patchify(torch.randn(4, 3, 224, 160)).shape)
print(patchify(torch.randn(1, 3, 32, 64)).shape)
patchify(torch.randn(1, 3, 30, 64))

torch.Size([4, 140, 768])
torch.Size([1, 8, 768])
RuntimeError: Expected input *args[0].shape[2] = 30 to be of the form 16*s68, where s68 is an integer

The batch b stayed dynamic, and so did the image size, as a multiple of 16: 16gh became the derived dimension 16 * gh, and the program refuses a height of 30. You wrote none of that. Each generic axis becomes a torch.export.Dim, shared wherever it appears; dims and literal sizes are fixed.

--aoti writes an AOTInductor package instead, compiled ahead of time, which C++ can load too (it needs a C++ compiler). From Python, shapescript.export("patches.ss", "patchify", "patchify.pt2") does the same and returns the ExportedProgram.

What can be exported

torch.export records torch operations, so an exported function, and everything it calls, can’t use Python. The checker enforces it as you type:

import python torch

dim D = 768

export fn load_qkv() -> [D 3*D]:
    return torch.load("w_qkv.pt") as [D 3*D]

$ shapescript check export_bad.ss
error[E1701]: `torch.load` is Python, which an exported function can't use
 --> export_bad.ss:6:12
  |
6 |     return torch.load("w_qkv.pt") as [D 3*D]
  |            ^^^^^^^^^^^^^^^^^^^^^^ runs Python
  |
  = note: `torch.export` records torch operations only, and code from `import python` could be anything
  = help: pass the value in as a parameter, or declare a torch function with `extern fn`, as in `extern fn erfc(x: Real) -> NonNeg = "torch.special.erfc"`

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

Also not allowed: models, modes, datasets, logs, gradients and optimizers, and the built-ins that use them. An exported program is for inference; training happens before it. An extern fn may be used if it names a function in torch (or math), so give a torch function a type with extern fn rather than calling it through import python.

What does carry over:

• as assertions become facts in the program. rows as [196 D] tells torch.export the size, and the program checks it when it runs.
• Unknown sizes (?k) become data-dependent sizes.
• if and for become torch.cond and torch.while_loop where they depend on values or dynamic sizes, and plain Python, unrolled, where they don’t.
• Iterator chains become torch.vmap or plain tensor operations, which torch.export records.

Serving a language model

A language model is served differently from how it’s trained. A serving engine such as vLLM or llama.cpp takes requests from many users at once, packs their tokens one after another along the sequence axis, and keeps the keys and values of past tokens in a cache, so each new token costs one step rather than a whole forward pass. To do that, an engine runs three parts of the model itself: the token embedding, attention and the output layer. Usually someone then writes the whole model again for each engine, and hopes it computes the same thing.

shapescript export --hf writes the model for vLLM instead, with everything but those three parts compiled from your source. Here is a small language model, which we’ll make servable:

dim D  = 64
dim H  = 4
dim Dh = D / H
dim V  = 500
dim L  = 2

model Tiny:
    param wte:   [V D]     = normal(std=0.02)
    param w_qkv: [L 3*D D] = normal(std=0.02)
    param w_o:   [L D D]   = normal(std=0.02)
    param norm:  [L D]     = ones

fn attention(x: [b n D], w_qkv: [3*D D], w_o: [D D]) -> [b n D]:
    q, k, v = linear(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))
        |> linear(w_o)

fn forward(m: Tiny, ids: [b n] Index[V]) -> [b n V] Logit over V:
    x = embedding(ids, m.wte)
    for l in 0..L:
        x = x + attention(rms_norm(x, over=D, weight=m.norm[l]), m.w_qkv[l], m.w_o[l])
    return linear(x, m.wte)

It checks, and it trains. shapescript export tiny.ss --hf tiny_hf writes it as a Hugging Face folder, which transformers loads like any other model; any model can be exported that way, a classifier or a speech recognizer too, with its function as the model’s forward. But it notes that vLLM can’t serve this one, and --vllm, which asks for a model vLLM can serve, says why:

$ shapescript export tiny.ss --hf tiny_hf --vllm
error[E3304]: a serving engine can't run this: this works along `n`, which mixes its positions
  --> tiny.ss:18:12
   |
18 |         |> softmax(over=m)
   |            ^^^^^^^ works across `n`
   |
   = note: a serving engine packs the tokens of many requests along `n`, one request after another, so everything but attention must treat each position on its own; it replaces `scaled_dot_product_attention(…, attn_mask=causal[n n])` with its own, which keeps the requests apart
   = help: write attention as `scaled_dot_product_attention` with `attn_mask=causal[n n]` in the call, and take positions as a parameter `pos: [b n] Index[P]` rather than computing them from `n`

With requests packed one after another, a softmax over every position in the batch would mix one user’s tokens into another’s. The engine’s own attention keeps them apart, but it can only stand in for attention it recognizes. Attention written by hand, with einsum and softmax, could be anything. The checker follows n through every function forward passes it to, and the only thing allowed to work across positions is scaled_dot_product_attention with attn_mask=causal[n n] in the call. So write attention with it:

fn attention(x: [b n D], w_qkv: [3*D D], w_o: [D D]) -> [b n D]:
    q, k, v = linear(x, w_qkv) |> split(b n (3 H Dh) -> b H n Dh)
    return scaled_dot_product_attention(q, k, v, attn_mask=causal[n n])
        |> rearrange(b H n Dh -> b n (H Dh))
        |> linear(w_o)

It computes the same thing, and now the export succeeds:

$ shapescript export tiny.ss --hf tiny_hf --vllm --weights ckpt.pt
wrote tiny_hf: 2 attention layers of 4 heads, vocabulary of 500; vLLM can serve it

tiny_hf is an ordinary Hugging Face folder. Its modeling_shapescript.py is your forward, compiled, with three changes: the embedding and the output layer are left to the engine, and each attention calls whichever attention function the engine sets. transformers uses your attention by default, so it gives your model’s logits:

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("tiny_hf", trust_remote_code=True)

and vLLM uses its own paged attention, with everything else running as your code says:

vllm serve tiny_hf --trust-remote-code --model-impl transformers --skip-tokenizer-init

There is nothing to keep in sync: change the model, export again.

The other rules are like this one. forward must take the model, the token ids and optionally their positions; embed the ids once, with embedding; and end by returning linear(h, m.W), the output layer, which the engine runs for the tokens it samples. A model that needs positions takes them as a parameter, pos: [b n] Index[P], rather than computing them from n: in a packed batch, the index along n isn’t a token’s position.

llama.cpp is different: a GGUF file holds weights and the name of an architecture, and llama.cpp runs its own implementation of that architecture. So --gguf doesn’t compile your model. It works out which of llama.cpp’s architectures your model is, by computing both on random weights, and refuses if none computes what yours does:

$ shapescript export tiny.ss --gguf tiny.gguf
error: `Tiny` isn't an architecture llama.cpp has: not `llama`: it has 3 fields stacked along its 2 blocks and 0 other fields, but `llama` has 9 tensors in each block and 1 other tensor; not `gpt2`: llama.cpp's architecture normalizes with `layer_norm` only, but the model uses rms_norm
error: exporting `forward` failed: `python3` exited with exit status: 1

The repository’s examples/llama.ss is a model --gguf writes as llama, and examples/gpt2.ss one it writes as gpt2.

Exercises

1. Export causal_attention from the Introduction. Which of its axes are dynamic in the program, and which are fixed?

2. Why is this an error, and how do you fix it while keeping it exportable?

   import python torch
   
   dim D = 64
   
   export fn act(x: [b D]) -> [b D]:
       return x |> torch.nn.functional.mish
   

3. Give Tiny learned position embeddings: a param wpe: [P D], and a parameter pos: [b n] Index[P] of forward whose embeddings are added to the tokens’. Why must the positions be a parameter, rather than computed with n |> map(i => i)? What does shapescript export --hf --vllm say if you try?

4. Make Tiny’s blocks those of examples/llama.ss: an RMSNorm and a SwiGLU feed-forward layer after attention, rotary embeddings, and a final norm. Export it with --gguf, and serve it with llama-server.

Several Devices

Every program so far began with devices cpu. This chapter is about the other choices: a GPU, several GPUs, and a GPU rented by the second on Modal, with no change to the code but that line.

One device

devices says where every tensor the program makes lives. There’s no default, and no code below that line mentions a device:

devices cpu          # the CPU
devices mps          # an Apple GPU
devices cuda[0]      # the first NVIDIA GPU

A tensor passed in from Python must already be on that device; it isn’t moved for you, and one on another device raises shapescript.DeviceError, naming the argument and suggesting x.to("cuda:0"). Datasets copy their batches to the device themselves, and logs and print copy values back.

Several devices

Change the line to several devices and the same program trains on all of them, with data parallelism:

devices cuda[0 1 2]

dim D  = 784
dim Hd = 256
dim C  = 10

model Mlp:
    param w1: [D Hd] = normal(std=0.02)
    param w2: [Hd C] = normal(std=0.02)

data Digit:
    img:   [D]
    label: [] Index[C]

fn loss(m: Mlp, batch: [b] Digit) -> [] NonNeg:
    return cross_entropy(batch.img @ m.w1 |> gelu |> matmul(m.w2), batch.label, over=C)

fn step(mut m: Mlp, mut opt: AdamW[Mlp], batch: [b] Digit) -> [] NonNeg:
    l, g = value_and_grad(loss, m, batch)
    update(mut m, mut opt, g)
    return l

fn train(mut m: Mlp, data: Dataset[Digit], epochs: int):
    opt = adamw(m, lr=3e-4)
    for epoch in 0..epochs:
        for batch in data |> shuffle(seed=epoch) |> batches(64):
            print("loss", step(mut m, mut opt, batch))

How it works, and what it promises:

• Your program runs once. One process, the controller, runs it, so print, save and calls into Python happen once. Each device gets a worker process that runs only the compiled tensor operations. You never ask which GPU you’re on.
• Device work and host work are separate functions. A function that loops over a dataset, prints, saves or calls Python runs in the controller; value_and_grad, update and the maths run on the workers. So step above is its own function, which train calls: putting value_and_grad directly in train’s loop is error E1803.
• The results are one device’s, up to the order floating-point numbers are added. A batch of 64 is split 22/21/21, the gradients are combined weighted by each part’s size, and the logged loss is the whole batch’s. batch_norm’s statistics cover the whole batch.
• The cost is visible. In the editor, hovering a line shows how each value is split across the devices and how many bytes each step sends between them: on this example, value_and_grad all-reduces the gradients of w1 and w2, 203,264 f32 values, 813,056 bytes per step.

Operations across the batch that can’t be split exactly, such as a softmax over the batch axis, are compile errors for now.

What works today. All of this has run on three GPUs on Modal, 3×T4 and 3×A100, with NCCL: the program runs once from one process, print prints once, and 100 training steps match a devices cuda[0] run to 2.4e-7 with dropout off. shard params and split H have run there too. Each training step is recorded as a CUDA graph and replayed, on each worker as on one GPU, which the fixed shapes allow (every step has the same shapes, with the last batch padded and masked). On 3×A100 a step of the training example takes 1.37 ms at batch 64, under PyTorch’s DDP on the same machine (2.37 ms), when its data is a TensorDataset that stays on the GPUs; a dataset that loads each batch on the CPU is slower with large batches. Without GPUs, devices cpu[0 1 2] runs a program just as devices cuda[0 1 2] would, with three worker processes on the CPU in place of three GPUs, so you can test a several-GPU program on a laptop. With dropout on, each device draws masks for its own examples, so results differ from one device’s. Coming: training across several machines.

Options

Indented lines under devices change how the work is split. They change speed and memory, never results:

devices cuda[0 1 2]
    shard params    # split parameters and optimizer state across the GPUs, for models too big for one
    split H         # split every parameter axis named H across the GPUs, such as attention heads

dim H = 12

With shard params, each device holds a third of every parameter and of its optimizer state, as FSDP does. With split H, every value with an axis of size H is split along it, so each device computes its own attention heads. H must be a dim that the number of devices divides.

Python on split data

Python code runs once, in the program’s own process, so it can’t run on data that’s split across devices. The checker says so:

devices cuda[0 1 2]

dim D = 784

data Digit:
    img: [D]

import python augment

fn augmented(batch: [b] Digit) -> [b D]:
    return augment.random_crop(batch.img) as [b D]

$ shapescript check split_python.ss
error[E1901]: `batch.img` is split across the 3 devices, but `augment.random_crop` is Python, which runs in the controller
  --> split_python.ss:11:32
   |
11 |     return augment.random_crop(batch.img) as [b D]
   |            ------------------- ^^^^^^^^^ split along its batch axis `b`
   |            |
   |            a Python function
   |
   = note: with 3 devices each one holds part of every batch, and Python calls run once, in the controller; this one would gather the batch through the controller's CPU at every step
   = help: move `random_crop` into the dataset, so it runs on each example as it's loaded, or declare it per example with `extern fn random_crop(...) = "augment.random_crop"`, so each device runs it on its own examples

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

An extern fn works on each example independently, so each device’s worker runs it on its own examples:

devices cuda[0 1 2]

dim D = 784

data Digit:
    img: [D]

extern fn random_crop(img: [b D]) -> [b D] = "augment.random_crop"

fn cropped(batch: [b] Digit) -> [b D]:
    return random_crop(batch.img)

Renting a GPU on Modal

No NVIDIA GPU at hand? Modal rents them by the second, and shapescript run --modal runs a whole program there, unchanged:

pip install modal          # into the Python that has shapescript
modal setup                # once: sign in
shapescript run --modal T4 train.py 3

The program’s folder is uploaded, the program runs in a container with that GPU and the same versions of Python, torch and ShapeScript, its output is streamed back as it’s printed, and the files it writes, such as checkpoints and logs, come back when it ends. In the container, devices cuda[0] is the GPU; --modal H100:3 gives three, for devices cuda[0 1 2]. --timeout 10m caps a run, --env KEY passes an environment variable, and --secret NAME a Modal secret. A .ss program is checked here first, so its errors cost nothing. The guide is docs/modal.md.

Exercises

1. You have a laptop and want to test that a program meant for devices cuda[0 1 2] gives the same losses as on one device. What do you change?

2. Why does devices cuda[0 1 2] with --modal T4 fail before anything starts, and what’s the fix?

3. Does this check?

   devices cuda[0 1 2]
       split H
   
   dim H = 8
   

Final Project: a Small Transformer

Time to put everything together. In this chapter we’ll build a small decoder-only transformer, a GPT in miniature, and train it end to end on a toy task that it can solve perfectly: copying. Each sequence is 8 random tokens followed by the same 8 again. The first half can’t be predicted, but once the model has seen it, every token of the second half can be, if the model learns to look back 8 positions. That’s exactly what attention is for.

The whole thing trains on a laptop CPU in a few seconds.

The plan

Three files in one folder:

• copier.ss: the sizes, the model, the data record, and the forward pass and loss;
• train.ss: training with the standard library’s fit;
• run_copier.py: the dataset, in Python, and the program that runs it all.

Sizes

devices cpu

## Vocabulary size: the tokens are 0 to 15.
dim V   = 16
## Length of each sequence: 8 random tokens, then the same 8 again.
dim T   = 16
## Model width.
dim D   = 64
## Number of attention heads.
dim H   = 4
## Width of each head.
dim Dh  = D / H
## Width of the MLP's hidden layer.
dim Dff = 4 * D
## Number of decoder blocks.
dim L   = 2

Everything that follows is checked against these. Change H to 5 and Dh is an error at once, since 5 doesn’t divide 64.

The model

A GPT’s parameters are a token embedding, a position embedding, and for each block, two layer norms, the attention projections and an MLP. We stack the blocks’ weights along a first axis of size L, so the model is one record:

## A small decoder-only transformer.
model Copier:
    param wte:   [V D]     = normal(std=0.02)
    param wpe:   [T D]     = normal(std=0.02)
    param ln1_w: [L D]     = ones
    param ln1_b: [L D]     = zeros
    param w_qkv: [L 3*D D] = normal(std=0.02)
    param w_o:   [L D D]   = normal(std=0.02)
    param ln2_w: [L D]     = ones
    param ln2_b: [L D]     = zeros
    param w1:    [L Dff D] = normal(std=0.02)
    param w2:    [L D Dff] = normal(std=0.02)
    param lnf_w: [D]       = ones
    param lnf_b: [D]       = zeros

linear(x, w) multiplies by the transpose of w, as in PyTorch, which is why w_qkv is [3*D D], output first.

The data

One example is a sequence and, for each position, the token that comes next:

## One training example: a sequence, and each position's next token.
data Example:
    ids:  [T] Index[V]
    next: [T] Index[V]

Attention

This is the function from the Introduction, with linear for the projections:

## Self-attention in which each token sees itself and the tokens before it.
fn attention(x: [b n D], w_qkv: [3*D D], w_o: [D D]) -> [b n D]:
    q, k, v = linear(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))
        |> linear(w_o)

Every line is checked: that 3 * H * Dh is 3 * D, that causal[n m] is a [n m] mask that broadcasts by name to [b H n m], that softmax normalizes over the keys m, and that the result is [b n D] again.

A block, and the forward pass

A pre-norm decoder block adds attention and then an MLP to its input. m.w_qkv[l] picks block l’s weights:

## Decoder block `l`: attention, then an MLP, each added to its input.
fn block(m: Copier, x: [b n D], l: int) -> [b n D]:
    h = x + (x
        |> layer_norm(over=D, weight=m.ln1_w[l], bias=m.ln1_b[l])
        |> attention(m.w_qkv[l], m.w_o[l]))
    return h + (h
        |> layer_norm(over=D, weight=m.ln2_w[l], bias=m.ln2_b[l])
        |> linear(m.w1[l])
        |> gelu
        |> linear(m.w2[l]))

## The logits of each position's next token.
fn forward(m: Copier, ids: [b T] Index[V]) -> [b T V] Logit over V:
    x = embedding(ids, m.wte) + m.wpe
    for l in 0..L:
        x = block(m, x, l)
    return x
        |> layer_norm(over=D, weight=m.lnf_w, bias=m.lnf_b)
        |> linear(m.wte)

embedding(ids, m.wte) + m.wpe adds a [T D] to a [b T D], which lines up by name: no unsqueeze. The loop runs the L blocks; x keeps its type each time round, as a loop requires. The output layer reuses the token embedding, linear(m.wte), a common trick called weight tying.

Loss and predictions

fn loss(m: Copier, batch: [b] Example, mode: Mode) -> [] NonNeg:
    return cross_entropy(forward(m, batch.ids), batch.next, over=V)

## The predicted next tokens.
fn predict(m: Copier, ids: [b T] Index[V]) -> [b T] Index[V]:
    return argmax(forward(m, ids), over=V)

loss takes a Mode, because that’s the signature the standard library’s fit expects. This model has no dropout, so it doesn’t use it.

Check the file:

shapescript check copier.ss

and hover Copier in the editor: it draws the architecture, with the embedding, the transformer block × L = 2 with its attention and MLP, and the tied output layer.

Training

train.ss imports the model and hands it to the standard library’s generic fit, which uses our loss:

import copier
import shapescript.std.fit as std

## Trains the copier with the standard library's loop: AdamW, batches of
## 64, a checkpoint after every epoch.
fn train(mut m: copier.Copier, data: Dataset[copier.Example], sink: LogSink, epochs: int):
    std.fit(mut m, data, sink, epochs, 64, 3e-3, Eval)

The program

The dataset is a few lines of PyTorch, in run_copier.py:

import sys

import torch

import shapescript
from copier import Copier, predict
from train import train


class CopyTask:
    """`n` sequences of 8 random tokens followed by the same 8 again."""

    def __init__(self, n, seed):
        g = torch.Generator().manual_seed(seed)
        half = torch.randint(0, 16, (n, 8), generator=g)
        seqs = torch.cat([half, half], dim=1)              # (n, 16)
        self.ids = seqs
        self.next = torch.roll(seqs, -1, dims=1)           # each position's next token

    def __len__(self):
        return len(self.ids)

    def __getitem__(self, i):
        return {"ids": self.ids[i], "next": self.next[i]}


epochs = int(sys.argv[1]) if len(sys.argv) > 1 else 10
shapescript.loader_workers(0)        # the dataset is defined in this script
model = Copier.init(seed=0)
sink = shapescript.sinks.csv("log.csv", every=None)  # a row per epoch
train(model, CopyTask(4096, seed=0), sink, epochs)
sink.close()

test = CopyTask(256, seed=1)
with torch.no_grad():
    guess = predict(model, test.ids)
# Positions 8 to 14 predict tokens 9 to 15, copies of tokens 1 to 7.
right = (guess[:, 8:15] == test.next[:, 8:15]).float().mean().item()
print(f"second-half accuracy: {right:.1%}")

Run it:

$ shapescript run run_copier.py 10
epoch 0 done
epoch 1 done
epoch 2 done
epoch 3 done
epoch 4 done
epoch 5 done
epoch 6 done
epoch 7 done
epoch 8 done
epoch 9 done
second-half accuracy: 100.0%
$ cat log.csv
loss,grad_norm
1.9936360102146864,0.924981951713562
1.2352662440389395,0.38128986954689026
1.223402040079236,0.35159793496131897
1.220885556191206,0.31604573130607605
1.2193415965884924,0.31824299693107605
1.2186482679098845,0.6085124611854553
1.2172029335051775,0.6178174614906311
1.2268130257725716,1.7973766326904297
1.2175495941191912,0.659481406211853
1.2164913676679134,0.2807769179344177

The model copies perfectly. Why does the loss stop at about 1.22 rather than 0? Of the 16 next tokens in each sequence, 7 are random (the second to eighth tokens), and no model can predict them better than chance, ln 16 ≈ 2.77 each. The other 9 are copies. So the best possible mean loss is 7/16 × 2.77 ≈ 1.21, which is where it ended.

fit also left a checkpoint, ckpt.pt, which std.resume can continue from. To train on a GPU, change devices cpu to devices cuda[0], or run it on Modal as it is, with devices cuda[0]: shapescript run --modal T4 run_copier.py 10.

Where to go from here

You’ve now used nearly all of ShapeScript. The language reference has the exact rules, and a section for every diagnostic; examples/ in the repository has more programs, such as a ViT block and a GPT. The exercises below extend the project.

Exercises

1. Add dropout after the MLP’s gelu, with p = 0.1. What has to change in block, forward and loss, and what does train pass now?
2. Change the task to reversal: the second half is the first half reversed. Which file changes? Will the model still reach 100%?
3. Write accuracy(m: Copier, batch: [b] Example) -> [] in ShapeScript: the fraction of positions 8 to 14 predicted correctly.
4. Export predict with shapescript export. What has to change first?

Appendix A: Diagnostics

Every error, warning and hint the checker reports has a code. Errors start with E, warnings with W and hints with H; the first two digits say which part of the language it’s about. This appendix lists them all, each linked to its section in the language reference, which explains it with a short program that produces it.

From the command line, shapescript explain E0401 prints the same section. In the editor, each diagnostic’s code is a link to it.

The ones you’ll meet most are E0401 (sizes that don’t match), E0403 (the wrong number of axes), E0404 (axes that don’t line up by name) and E0501 (a value that doesn’t fit a domain).

Lexical structure

Syntax

Names and declarations

Shapes and sizes

Value domains and dtypes

Math in comments

Notebooks

Control flow and host scalars

Models

Devices

mut and gradients

Modes

Data and datasets

Logs and sinks

Optimizers

Checkpoints

Export

Several devices

Options for several devices

Fixed shapes and padding

Grouped sizes and shape aliases

Modules and generic functions

Iterator chains: what they iterate and their independence

Python modules for .ss files

Packages

Iterator chains: syntax and stages

Appendix B: Glossary

Axis. One entry of a tensor’s shape, such as n in [b n D]. An axis has a size and may have a name.

as. An assertion: x as T claims x already has type T, which the checker then believes and the program checks at run time. The only way to tell the checker something it can’t prove. See Chapter 7.

Batch. [b] Digit: b examples of a data record, whose fields each get the axis b first.

Broadcasting by name. How elementwise operations line up tensors: by axis name when one operand’s names are all in the other, otherwise from the right, where two named axes in a pair must have the same name. See Chapter 6.

Buffer. A model field that’s saved but not trained.

Controller. With several devices, the one process that runs your program; the workers run its tensor operations.

Data record. A data declaration: the fields of one example of a dataset.

Devices. The devices line: where a program’s tensors live, such as cpu, mps, cuda[0], or several, cuda[0 1 2].

Dim. A named constant size, dim D = 768. Writing D in a type both sets the axis’s size and names it.

Domain. What values a tensor may hold, such as Unit, Prob over C or Index[V]. The checker infers most domains from the code. See Chapter 8.

Dtype. How values are stored: f32, bf16, u8, bool and so on.

Elementwise function. A function whose parameters have no shape, such as fn swish(x: Real) -> Real, which applies to tensors of any shape.

Export. Turning a function into a torch.export program that runs without ShapeScript or Python code. See Chapter 15.

extern fn. A typed declaration of a Python function, whose signature the checker trusts.

Generic axis. A name in a parameter’s type that isn’t declared anywhere, such as b in x: [b D]. Each call chooses its size.

Generic function. A function with a type variable, such as M in m: M, checked again for the types each call gives it.

Grad. Grad[M], a gradient with respect to a model M’s params, with a field for each.

Group. (gh 16): one axis whose size is the product of its parts.

Grouped size. 16gh: shorthand for the group (gh 16), an axis whose size is a multiple of 16.

Host scalar. An int, float or bool: a plain Python value, kept on the host, for counters, rates and flags.

Logs record. A logs declaration: named metrics and how each combines.

Mask. A boolean tensor, from comparisons or causal[n m].

Mode. Train(seed=s) or Eval, passed to everything random, such as dropout.

Model. A model declaration: a record of params and buffers, the state of a network.

mut. Marks a parameter a function may change, and the argument at every call.

Param. A model field that’s trained.

Pattern. b n (h d) -> b h n d: the einops-style description an operation such as rearrange or einsum takes, written as syntax.

Pipe. a |> f(args), which means f(a, args).

Prelude. The extern fn declarations every file sees, such as exp, gelu and layer_norm.

Python value. A value from import python, which the checker knows nothing about until an as gives it a type.

Schedule. A learning rate that changes with the step, such as warmup(100, cosine(3e-4, 10000)).

Shape alias. dim Images = [b 3 16gh 16gw]: a name for a list of axes.

Sink. A Python object that logs records are written to, such as shapescript.sinks.csv("log.csv").

to. A conversion, such as u8 to Unit or a mask to Additive. Unlike as, it changes the values.

Type variable. A single capital letter, optionally followed by digits, in a parameter’s type that nothing declares: M, E.

Unknown size. ?k: a definite size the checker can’t see, such as the number of rows a filter keeps.

Appendix C: From PyTorch to ShapeScript

A cheat-sheet for translating PyTorch code. The ShapeScript column assumes the types on the left: x: [b n D] and so on.

Shapes and reshaping

Reductions and operations along an axis

Products

Masks and attention

Dtypes and conversions

Models and training

One function, both ways

PyTorch:

def attention(x, w_qkv, w_o, n_heads):
    B, N, D = x.shape
    q, k, v = (x @ w_qkv).view(B, N, 3, n_heads, D // n_heads).permute(2, 0, 3, 1, 4)
    att = (q @ k.transpose(-2, -1)) / math.sqrt(D // n_heads)
    mask = torch.tril(torch.ones(N, N, dtype=torch.bool, device=x.device))
    att = att.masked_fill(~mask, float("-inf")).softmax(dim=-1)
    return (att @ v).transpose(1, 2).reshape(B, N, D) @ w_o

ShapeScript:

dim D  = 768
dim H  = 12
dim Dh = D / H

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

The PyTorch version has a view, a permute, a transpose and a reshape, each of which could be wrong while still running. The ShapeScript version has the same steps, each checked.

Appendix D: Setup Troubleshooting

Most setup problems come down to one question: which Python, and which shapescript, is running? So start by asking, in the terminal where the problem happens:

python -c "import sys; print(sys.executable)"
which -a python shapescript

python -c "import sys; print(sys.executable)"
type -a python shapescript

python -c "import sys; print(sys.executable)"
Get-Command python, shapescript -All

With the environment active, the first one listed is in your project’s .venv. The rest of this appendix is the usual problems, roughly in the order you’d meet them in Chapter 1.

Building from source

These come up only if you build ShapeScript yourself (Chapter 1, Building from source).

linker 'cc' not found, or link.exe not found

Rust needs a C linker. On macOS, xcode-select --install; on Debian and Ubuntu, sudo apt install build-essential; on Fedora, sudo dnf install gcc; on Windows, Microsoft’s C++ build tools (Chapter 1 has the winget command).

rustc 1.xx is not supported

ShapeScript needs Rust 1.88 or later. Run rustup update.

cargo: command not found

Open a new terminal: rustup changed the PATH of new ones. In fish, run fish_add_path ~/.cargo/bin once.

Python and the virtual environment

Windows: Activate.ps1 cannot be loaded

PowerShell refuses to run scripts by default, and activating is a script. Allow your own and signed ones, for your user only:

Set-ExecutionPolicy -Scope CurrentUser RemoteSigned

Then activate again.

Windows: python opens the Microsoft Store

That’s Windows’ placeholder for a Python that isn’t installed, or isn’t on the PATH. Use the launcher, py -3.13, to make the environment; once it’s active, python is the environment’s.

python isn’t the environment’s, even though it’s active

Something put another Python ahead of .venv on your PATH: usually a version manager such as pyenv, asdf or mise that switches Python when you change folders, or conda’s base environment. Check with the commands at the top. Either turn the tool off for the project, or skip activation and name the environment’s programs by path: .venv/bin/python -m pip install ... (.venv\Scripts\python.exe on Windows).

Apple Silicon: Bad CPU type in executable, or incompatible architecture

The Python is an Intel (x86_64) one, and the programs you built are for Apple Silicon (arm64), or the other way round. Errors like mach-o file, but is an incompatible architecture (have 'arm64', need 'x86_64') on import shapescript, or pip finding no version of torch at all, are the same problem. Check:

python -c "import platform; print(platform.machine())"
uname -m

Both should print arm64. If Python prints x86_64, it’s an Intel Python, often from an old Homebrew in /usr/local: delete .venv and make it again with /opt/homebrew/bin/python3.13 -m venv .venv. If uname -m prints x86_64, your terminal runs under Rosetta: in Finder, open the terminal app’s Get Info and untick Open using Rosetta.

pip install shapescript: No matching distribution found

pip found no wheel of ShapeScript for this Python and platform. The wheels are for Python 3.9 or later on Linux (x86_64 and aarch64, with glibc 2.28 or later), macOS (Apple Silicon on macOS 11 or later, Intel on 10.12 or later) and Windows (x64). Check the Python with python --version and python -c "import platform; print(platform.machine())": an old Python, a 32-bit one, or an Intel Python on Apple Silicon is the usual cause (see the entry above). An old pip can miss the wheels too: python -m pip install --upgrade pip. On another platform, such as Windows on Arm or Alpine Linux, build from source.

maturin: Both VIRTUAL_ENV and CONDA_PREFIX are set

conda’s base environment is active underneath your virtual environment, and maturin develop doesn’t know which to install into. Leave conda, then activate .venv again:

conda deactivate

Or, for this terminal only, unset conda’s variable:

unset CONDA_PREFIX

set -e CONDA_PREFIX

Remove-Item Env:CONDA_PREFIX

conda config --set auto_activate_base false stops conda from activating base in every new terminal. Or use pip install ~/shapescript/crates/py, which doesn’t mind conda.

this installation of shapescript has no shapescript binary

The package was built from source without the two programs. Copy them into the package (Chapter 1, Clone and build) and run pip install ~/shapescript/crates/py again. With maturin develop, the copy alone is enough.

shapescript: command not found

The environment isn’t active in this terminal, or ShapeScript isn’t installed in it. Activate it, and check with pip show shapescript.

shapescript run can’t find torch or shapescript

If shapescript run says No module named shapescript, or can’t import torch, it ran the wrong Python. shapescript run and shapescript export run Python: the environment’s own shapescript uses the environment’s Python; one installed with cargo install uses --python, else the SHAPESCRIPT_PYTHON environment variable, else python3. Pass --python .venv/bin/python, or set the variable:

export SHAPESCRIPT_PYTHON=~/hello-shapescript/.venv/bin/python

set -x SHAPESCRIPT_PYTHON ~/hello-shapescript/.venv/bin/python

$env:SHAPESCRIPT_PYTHON = "$HOME\hello-shapescript\.venv\Scripts\python.exe"

These last for the terminal. (In fish, set -Ux keeps it for good, and set -e SHAPESCRIPT_PYTHON removes it.)

pip install ./mypkg is slow, or fails without a network

pip builds a package in a fresh environment, with its build dependencies from PyPI: ShapeScript, and PyTorch with it, which is a large download, and impossible offline. Build it with your environment’s ShapeScript instead:

pip install hatchling
pip install --no-build-isolation ./mypkg

PyTorch

torch.cuda.is_available() is False on a machine with an NVIDIA GPU

Look at torch.__version__. If it ends in +cpu, you have the CPU build, the default on Windows: uninstall it and install the CUDA build with the command from pytorch.org’s selector:

pip uninstall -y torch

If PyTorch says The NVIDIA driver on your system is too old, the build needs a newer CUDA than your driver supports. nvidia-smi shows the newest CUDA version the driver supports, top right: pick a build no newer than that, or update the driver. In a notebook, restart the kernel after changing PyTorch.

UserWarning: Failed to initialize NumPy

PyTorch works without NumPy, but warns. pip install numpy.

The editor

The editor shows no errors or hovers

Go down this list:

1. The status bar, bottom right, should say ShapeScript when a .ss file is open. If it says Plain Text, the extension isn’t installed in this VS Code: code --list-extensions should list shapescript.shapescript. Reload the window after installing it.
2. If VS Code opened the folder in Restricted Mode, trust it (the banner at the top, or Workspaces: Manage Workspace Trust).
3. Open the Output panel (View › Output) and choose ShapeScript in its menu. An error there, or a notification saying couldn't start the language server, usually means a shapescript.server.path setting that names a file that doesn’t exist, or a .vsix for another platform (darwin-x64 on Apple Silicon, say). Fix or remove the setting, or uninstall the extension and install it again from the Marketplace, which picks your platform’s (or, if you build it, rebuild the .vsix for your platform).

ci/vsix.sh: warning: no dist/bin/…/shapescript-lsp, so no VSIX

The script found no language server for that platform. Build it first (cargo build --release -p shapescript-lsp), and name your own machine’s platform: the script packs target/release’s server only for the platform it was built on. On Windows, run the script in Git Bash.

Pylance: Import "torch" could not be resolved, or "shapescript"

VS Code is using another Python than your environment’s. Run Python: Select Interpreter and pick .venv, or set python.defaultInterpreterPath (Chapter 1, Select the environment’s Python). The status bar shows the selected one when a Python file is open.

Pylance: Import "hello" could not be resolved

There’s no hello.py beside hello.ss yet. The extension writes it when hello.ss is opened or saved without errors: fix its errors and save it. It writes none when the setting shapescript.pythonModules is off.

Debugging: debugpy isn’t found

The debugger comes with Microsoft’s Python extension. Without it, install it into the environment, pip install debugpy, and name the environment’s Python in the setting shapescript.debug.python.

Notebooks

%load_ext shapescript: No module named 'shapescript'

The kernel isn’t your environment’s Python. In VS Code, click the kernel’s name at the top right and pick .venv; in Jupyter in the browser, install jupyterlab into the environment and start it from there. import sys; sys.executable in a cell shows which Python the kernel is.

The kernel can’t start, or VS Code asks for ipykernel

The environment needs pip install ipykernel (VS Code offers to run it).