Keyboard shortcuts

Press ← or → to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

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

PyTorchShapeScript
x.shape # (B, N, D), in a commentx: [b n D], in the type
x.transpose(1, 2)rearrange(x, b n D -> b D n)
x.permute(0, 2, 1, 3)rearrange(x, b H n Dh -> b n H Dh)
x.view(B, N, H, D // H)rearrange(x, b n (H Dh) -> b n H Dh)
x.reshape(B, -1)rearrange(x, b n D -> b (n D))
x.unsqueeze(1)rearrange(x, b D -> b 1 D), or nothing: broadcasting is by name
x[:, None, :] + yx + y, when the names line up
x.expand(B, 4, D)repeat(x, b D -> b r D, r=4)
q, k, v = x.chunk(3, dim=-1)q, k, v = split(x, b n (3 D) -> b n D)
torch.cat([a, b], dim=1)cat(a, b, over=n)
torch.stack([a, b])stack(a, b, over=k)
x[:, 0], x[:, 1:], x[:, :-1]the same
x[mask]x[mask], of type [?k D]

Reductions and operations along an axis

PyTorchShapeScript
x.mean(dim=1)mean(x, over=n)
x.sum()sum(x)
x.softmax(dim=-1)softmax(x, over=D)
x.argmax(dim=-1)argmax(x, over=D), of type Index[D]
x.topk(5, dim=-1)vals, ids = topk(x, 5, over=D)
x.cumsum(dim=1)cumsum(x, over=n)
F.cross_entropy(logits, y)cross_entropy(logits, y, over=C)
F.layer_norm(x, (D,))layer_norm(x, over=D)
torch.gather(x, 1, ids)gather(x, ids, over=n)

Products

PyTorchShapeScript
x @ wx @ w
F.linear(x, w, b)linear(x, w, b)
torch.einsum("bhnd,bhmd->bhnm", q, k)einsum(q, k: b H n Dh, b H m Dh -> b H n m)

Masks and attention

PyTorchShapeScript
torch.tril(torch.ones(n, n)).bool()causal[n n]
scores.masked_fill(~mask, float("-inf"))masked_fill(scores, !mask, -inf)
scores + torch.where(mask, 0, -inf)scores + (mask to Additive)
F.scaled_dot_product_attention(q, k, v, is_causal=True)scaled_dot_product_attention(q, k, v, attn_mask=causal[n n])

Dtypes and conversions

PyTorchShapeScript
img.float() / 255img to Unit
x * 2 - 1 (for [0, 1] values)x to Signed
x.to(torch.bfloat16)x to bf16
mask.float()mask to Binary
assert x.shape == (196, D)x as [196 D]

Models and training

PyTorchShapeScript
class Mlp(nn.Module) with nn.Parametersmodel Mlp: with param fields
self.register_buffer("mean", ...)buffer mean: [D] = zeros
Mlp() with torch.manual_seed(0)Mlp.init(seed=0)
model.train() / model.eval()a Mode argument: Train(seed=s) / Eval
F.dropout(x, 0.1, training=...)dropout(x, 0.1, mode)
loss.backward(), p.gradl, g = value_and_grad(loss, m, ...)
optimizer.zero_grad()nothing: gradients are values
torch.optim.AdamW(model.parameters(), lr=3e-4)adamw(m, lr=3e-4), of type AdamW[Mlp]
clip_grad_norm_(params, 1.0); optimizer.step()update(mut m, mut opt, g |> clip_norm(1.0))
a LambdaLR schedulerlr=warmup(100, cosine(3e-4, 10000))
DataLoader(ds, batch_size=64, shuffle=True)ds |> shuffle(seed=epoch) |> batches(64)
torch.save({...}, "ckpt.pt")save("ckpt.pt", m, opt, epoch)
torch.load("ckpt.pt") and load_state_dictm, opt, done = load("ckpt.pt", Mlp, AdamW[Mlp], int)
x.to("cuda") everywheredevices cuda[0], once
DistributedDataParalleldevices cuda[0 1 2]

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.