ShapeScript

A typed tensor language that compiles to PyTorch. Axis names and sizes are part of the syntax, and the checker verifies every shape before any code runs. It's einops built into a language: no pattern strings, no dim=-1, no unsqueeze, and no dropping back into Python halfway through an expression.

dim D  = 768
dim H  = 12
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)

Swap two letters in the last einsum, writing b H m n for the attention weights, and PyTorch still runs it, since n and m have the same size. ShapeScript doesn't compile it: error E0408, the pattern calls axis 3 m, but m is axis 4.

This package is a pre-release. It has the compiler (shapescript), the language server (shapescript-lsp), the runtime, import of .ss files from Python and the notebook magic (%%shapescript), with prebuilt binaries for Linux (x86_64, aarch64), macOS (arm64, x86_64) and Windows (x64). It needs PyTorch:

pip install torch
pip install --pre shapescript
shapescript check model.ss

import shapescript
shapescript.install()   # then `import model` loads model.ss

The editor support is the ShapeScript extension for VS Code, on the Visual Studio Marketplace and Open VSX, with the language server inside it.

See the repository for the tour, and the language reference. Licensed under MIT or Apache-2.0, at your option.