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 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, with the language server inside it: code --install-extension shapescript.shapescript --pre-release. Other editors can run shapescript-lsp, which this package puts on the environment's PATH.

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