Metadata-Version: 2.4
Name: shapescript
Version: 0.1.4b1
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Rust
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Compilers
Requires-Dist: torch>=2.0
Requires-Dist: einops>=0.7
Requires-Dist: pytest ; extra == 'test'
Requires-Dist: numpy ; extra == 'test'
Provides-Extra: test
License-File: LICENSE-MIT
License-File: LICENSE-APACHE
Summary: Tensor programs whose shapes are checked before they run
Keywords: pytorch,tensor,shapes,einops,type checker,compiler
License-Expression: MIT OR Apache-2.0
Requires-Python: >=3.9
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
Project-URL: Documentation, https://gitlab.com/NatanFreeman/shapescript/-/blob/main/docs/reference.md
Project-URL: Homepage, https://gitlab.com/NatanFreeman/shapescript
Project-URL: Issues, https://gitlab.com/NatanFreeman/shapescript/-/issues
Project-URL: Repository, https://gitlab.com/NatanFreeman/shapescript

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

```sh
pip install torch
pip install shapescript
shapescript check model.ss
```

```python
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](https://marketplace.visualstudio.com/items?itemName=shapescript.shapescript), 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](https://gitlab.com/NatanFreeman/shapescript) for the tour, and the [language reference](https://gitlab.com/NatanFreeman/shapescript/-/blob/main/docs/reference.md). Licensed under MIT or Apache-2.0, at your option.

