einops, built into a language.

ShapeScript is a typed tensor language. Axis names and sizes are part of the syntax. Every shape is checked before any code runs. It compiles to Python that calls PyTorch, so you keep using the same libraries.

pip install --pre shapescript

VS Code extension

Read the book Source on GitLab

ShapeScriptattention.ss

## Model width.
dim D  = 768
## Number of attention heads.
dim H  = 12
## Width of each head. The checker verifies that D is divisible by H.
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‹b n 3*D›
        |> split(b n (3 H Dh) -> b H n Dh)‹b H n Dh› ×3

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

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

Causal attention. It compiles. The grey shapes at the line ends are what the checker infers, as the editor shows them. Each example below is code like this with one bug.

Six bugs PyTorch lets through

Each example is PyTorch code with a bug, and what happened when it ran. Under it is the same code in ShapeScript, with the error the checker reports.

PyTorchattention.py

def attention(x, w_qkv, w_o):
    q, k, v = (x @ w_qkv).view(b, n, 3, H, Dh).permute(2, 0, 3, 1, 4)
    att = torch.einsum("bhnd,bhmd->bhnm", q, k) / math.sqrt(Dh)
    att = att.masked_fill(~causal, float("-inf")).softmax(dim=-1)
    y = torch.einsum("bhmn,bhmd->bhnd", att, v)
    return y.transpose(1, 2).reshape(b, n, D) @ w_o

When runRuns. The output is [2, 16, 768] with no NaNs, and every token attends to the wrong positions: n and m are both the sequence length, so nothing can tell them apart.

ShapeScriptattention.ss

## Model width.
dim D  = 768
## Number of attention heads.
dim H  = 12
## Width of each head. The checker verifies that D is divisible by H.
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‹b n 3*D›
        |> split(b n (3 H Dh) -> b H n Dh)‹b H n Dh› ×3

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

    return einsum(att, v: b H m n, b H m Dh -> b H n Dh)error E0408 the pattern calls axis 3 m, but m is axis 4 of [b H n m] Prob over maxis 3 is npatterns name axes by positionto reorder the axes, use rearrange first        |> rearrange(b H n Dh -> b n (H Dh))
        |> matmul(w_o)

PyTorchloss.py

def loss(x, w, y):
    p = (x @ w).softmax(dim=-1)
    return F.cross_entropy(p, y)

When runRuns and returns a loss (6.91). cross_entropy applies its own softmax, so this one trains, just worse, and nothing says why.

ShapeScriptloss.ss

dim D = 768
dim C = 1000

fn loss(x: [b D], w: [D C], y: [b] Index[C]) -> [] NonNeg:
    p = x @ w |> softmax(over=C)‹b C› Prob over C
    return cross_entropy(p, y, over=C)error E0501 expected Logit over C, got Prob over Cthis is [b C] Prob over Csoftmax applied before cross_entropy?

PyTorchembed.py

def embed(emb, ids):
    return emb[ids]

embed(token_table, labels)    # labels are class ids in [0, 1000)

When runRuns. Every class id is a valid row of the 32,000-row token table, so the lookup succeeds and returns the wrong embeddings.

ShapeScriptembed.ss

dim V = 32000
dim D = 768
dim C = 1000

fn embed(emb: [V D], ids: [b t] Index[C]) -> [b t D]:
    return emb[ids]error E0507 expected Index[V], got Index[C]this is [b t] Index[C]the sizes differ: C = 1000, not V = 32000

PyTorchpool.py

def pool(x):              # x: (b, n, D)
    return x.mean(dim=-1)

When runRuns and returns (b, n) instead of (b, D). Nothing complains until a later layer does, and never if n happens to equal D.

ShapeScriptpool.ss

dim D = 768

fn pool(x: [b n D]) -> [b D]:
    return x |> reduce(b n D -> b n, mean)‹b n›error E0401 mismatched shape: expected [b D], found [b n]the return type is [b D]axis 2: n is not known to equal D = 768

PyTorchpatches.py

def forward(img, embed, head):
    x = img.unfold(2, 16, 16).unfold(3, 16, 16)
    x = x.permute(0, 2, 3, 4, 5, 1).reshape(img.shape[0], -1, 16 * 16 * 3)
    return (x @ embed).mean(dim=1) @ head

Fails at runtimeRuntimeError: mm: mat1 and mat2 shapes cannot be multiplied (8x768 and 512x192)

ShapeScriptpatches.ss

dim D = 192
dim C = 10
dim Images = [b 3 16gh 16gw]

model Classifier:
    param embed: [512 D] = normal(std=0.02)
    param head:  [D C]   = normal(std=0.02)

fn forward(m: Classifier, img: Images) -> [b C]:
    return img‹b 3 (gh 16) (gw 16)›
        |> rearrange(b c (gh p1) (gw p2) -> b (gh gw) (p1 p2 c))‹b (gh gw) (p1 p2 c)›
        |> matmul(m.embed)error E0401 can't multiply [b (gh gw) (p1 p2 c)] by [512 D]this is [512 D]the axes @ contracts differ: (p1 p2 c) = 16*16*3 = 768, not 512        |> reduce(b t D -> b D, mean)
        |> matmul(m.head)

PyTorchgate.py

def log_gate(x, w):
    g = torch.sigmoid(x @ w)
    return torch.log(g)

When runRuns, and returns -inf wherever sigmoid rounds to 0 (774 of 2,000 values in our test), which turns into NaN gradients later.

ShapeScriptgate.ss

dim D = 768
dim C = 1000

fn log_gate(x: [b D], w: [D C]) -> [b C]:
    g = x @ w |> sigmoid‹b C› Unit
    return log(g)warning W0501 g can be exactly 0log(0) is -infit's Unit, so its values are in [0, 1]use log_sigmoid(x @ w)

It compiles to Python

shapescript build writes ordinary Python that calls PyTorch. You can read it, import it from Python and use it in a notebook.

ShapeScriptattention.ss

## Model width.
dim D  = 768
## Number of attention heads.
dim H  = 12
## Width of each head. The checker verifies that D is divisible by H.
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‹b n 3*D›
        |> split(b n (3 H Dh) -> b H n Dh)‹b H n Dh› ×3

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

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

↓ shapescript build

Pythonattention.py

def causal_attention(x, w_qkv, w_o, _ss_pad=None):
    _ss_dev = _ss_rt.device()
    _ss_d_n = x.shape[1]
    _ss_pads = _ss_rt.pads(_ss_pad, {"n": 4})
    _ss_t1 = x @ w_qkv
    q, k, v = _ss_einops.rearrange(_ss_t1, "b n (i0 H Dh) -> i0 b H n Dh", i0=3, H=H, Dh=Dh).unbind(0)
    _ss_t2 = _ss_torch.einsum("bHnD,bHmD->bHnm", q, k)
    _ss_t3 = _ss_t2 * (1 / _ss_torch.sqrt(_ss_rt.to_float(Dh)))
    _ss_t4 = _ss_t3.masked_fill(_ss_torch.logical_not(_ss_rt.causal(_ss_d_n, _ss_d_n, _ss_dev)), -_ss_math.inf)
    att = _ss_rt.masked_along("softmax", _ss_t4, 3, _ss_pads, ((4, 3),))
    _ss_t5 = _ss_torch.einsum("bHnm,bHmD->bHnD", _ss_rt.fill(att, _ss_pads, ((4, 3),)), v)
    _ss_t6 = _ss_einops.rearrange(_ss_t5, "b H n Dh -> b n (H Dh)")
    return _ss_t6 @ w_o

shapescript build attention.ss --checks=off, the function only. Patterns become einops and einsum strings, D / H is worked out to 64, and _ss_rt is ShapeScript's small runtime.

ShapeScriptmlp.ss

dim D = 768

fn mlp(x: [b n D], w1: [D 4*D], w2: [4*D D]) -> [b n D]:
    return x @ w1 |> gelu |> matmul(w2)‹b n D›

↓ shapescript build

Pythonmlp.py

def mlp(x, w1, w2):
    _ss_dev = _ss_rt.device()
    return _ss_F.gelu(x @ w1) @ w2

shapescript build mlp.ss --checks=off, the function only. With checks on (the default), it also checks each argument's rank, dtype, device and sizes on entry.

Python and ShapeScript call each other

Each language calls the other directly, and one debugger steps through both. Your datasets, libraries and notebooks stay as they are.

Pythontrain.py

import shapescript
shapescript.install()                    # lets `import` find .ss files
shapescript.devices("cuda[0]")           # once per program, unless a .ss module declares `devices`

from model import causal_attention       # compiles model.ss (cached)
y = causal_attention(x, w_qkv, w_o)      # argument shapes are checked on entry

DeviceError: model.ss:8:21: argument `x` of `causal_attention` is on cpu,not on the declared device `cuda[0]`. Move it there first, as with `x.to("cuda:0")`

ShapeScriptmodel.ss

import python torch

dim D = 768

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

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

$ shapescript pymodules *.ss# write model.py beside model.ss

Stopped at a breakpoint on return y in causal_attention; the Variables view lists att = [b=2 H=4 n=10 m=10] float32, k, q and v = [b=2 H=4 n=10 Dh=16] float32, w_o = [D=64 D=64], w_qkv = [D=64 3*D=192], x = [b=2 n=10 D=64] and y, and the call stack shows causal_attention and main in attend.ss

Training without the footguns

You can still train from Python: a ShapeScript model is a torch.nn.Module, so your loop, Lightning or anything else works with it. But training in ShapeScript is better.

This is a whole training program, on three GPUs. Here are the mistakes PyTorch and Lightning allow that it can't make. Point at one to see the lines that rule it out.

Every tensor is made on the device the program declares. No code mentions a device.

ShapeScripttrain.ss

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]

logs StepLogs:
    loss:      [] mean
    grad_norm: [] max

fn forward(m: Mlp, x: [b D], mode: Mode) -> [b C] Logit over C:
    return x @ m.w1 |> gelu |> dropout(0.1, mode) |> matmul(m.w2)‹b C›

fn loss(m: Mlp, batch: [b] Digit, mode: Mode) -> [] NonNeg:
    return cross_entropy(forward(m, batch.img, mode), batch.label, over=C)

fn train_step(mut m: Mlp, mut opt: AdamW[Mlp], batch: [b] Digit, 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, grad_norm=n)

fn fit(mut m: Mlp, train: Dataset[Digit], sink: LogSink, epochs: int, mode: Mode):
    opt = adamw(m, lr=3e-4)
    run(mut m, mut opt, train, sink, 0, epochs, mode)

fn run(mut m: Mlp, mut opt: AdamW[Mlp], train: Dataset[Digit], 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")

The VS Code extension

The extension has the language server inside it. It shows the shape of every value as you write.

Hovering att shows att: [b H n m] Prob over m, lower-triangular, and a stack of squares, each with its lower triangle filled, labelled n and m, with × H × b beside the stack

Hovering the rearrange stage: [b H n Dh] on the left, an arrow, and [b n (H Dh)] on the right, its last axis drawn as H blocks of Dh

The classifier example in VS Code, with a shape hint at the end of each line of forward

gpt.ss on the left; on the right the Architecture tab for Gpt, from the embedding through a transformer block × 6 to a linear layer tied to wte, with the shape after each layer

Running and exporting

The same program exports for torch.export and Hugging Face, runs in notebooks, and runs on one GPU or several.

ShapeScriptpatchify.ss

dim D = 768

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))‹b (gh gw) (p1 p2 c)›

Pythonload.py

import torch
patchify = torch.export.load("patchify.pt2").module()    # no ShapeScript needed
tokens = patchify(torch.randn(4, 3, 224, 160))           # [4 140 768]

# any model, such as speech recognition: transformers loads the folder$ shapescript export asr.ss --fn transcribe --hf asr_hf # a language model: vLLM serves the same kind of folder$ shapescript export llama.ss --hf llama_hf --vllm$ vllm serve llama_hf --trust-remote-code --model-impl transformers --skip-tokenizer-init

For vLLM, the checker proves the model has the form a serving engine needs, and says which rule it breaks if it doesn't:

It takes token ids, [b n] Index[V], and optionally positions, and gives next-token logits, [b n V].
It starts with an embedding and ends with a linear output layer, which the engine runs itself.
Only causal scaled_dot_product_attention works across positions, so the engine can swap in its own attention and cache.

The engine runs the embedding, attention and output layer, and the rest is your source, compiled.

Notebook cellIn [2]

%%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)‹b D›

- devices cuda[0]+ devices cuda[0 1 2]

ShapeError: model.ss:13:14:expected [196 D] = (196, 768), got shape (195, 768)

Install the beta

The current release, 0.1.2b1, is a pre-release, so pip needs --pre.

Python package

The compiler, the language server, the runtime and the notebook magic.

pip install --pre shapescript

VS Code extension

Hovers, inlay hints and the Architecture tab. The language server is inside it.

code --install-extension shapescript.shapescript

Documentation

The book and the language reference are in the repository.

Read the book Reference Source on GitLab