Metadata-Version: 2.4
Name: stile-verifier
Version: 0.1.3
Summary: A type system for numerical programs: write a spec, get a structural proof your code computes it.
Project-URL: Homepage, https://github.com/save-buffer/stile
Project-URL: Repository, https://github.com/save-buffer/stile
Project-URL: Issues, https://github.com/save-buffer/stile/issues
Author-email: Sasha Krassovsky <krassovskysasha@gmail.com>
License-Expression: MIT
License-File: LICENSE
Keywords: einsum,jax,kernels,numerical,pallas,pytorch,triton,type-system,verification
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: Implementation :: CPython
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.11
Provides-Extra: all
Requires-Dist: einops>=0.8.1; extra == 'all'
Requires-Dist: jax>=0.4.0; extra == 'all'
Requires-Dist: numpy>=1.24.4; extra == 'all'
Requires-Dist: torch>=2.0.0; extra == 'all'
Requires-Dist: triton>=3.0.0; extra == 'all'
Provides-Extra: jax
Requires-Dist: einops>=0.8.1; extra == 'jax'
Requires-Dist: jax>=0.4.0; extra == 'jax'
Provides-Extra: numpy
Requires-Dist: numpy>=1.24.4; extra == 'numpy'
Provides-Extra: torch
Requires-Dist: einops>=0.8.1; extra == 'torch'
Requires-Dist: torch>=2.0.0; extra == 'torch'
Provides-Extra: triton
Requires-Dist: torch>=2.0.0; extra == 'triton'
Requires-Dist: triton>=3.0.0; extra == 'triton'
Description-Content-Type: text/markdown

# Stile: A Type System for Numerical Programs

**Formally Verify Your Numerical Programs**

Stile is a type system for numerical programs. Describe what your program should compute in a lightweight specification language and Stile will structurally prove the program's adherance to the spec.

A short demo of what this lets you write today:

```python
# Spec: causal flash attention, one line.
SPEC = (
    "(softmax[nctx where nctx <= qctx]"
    "((qctx dhead, nctx dhead -> qctx nctx) / sqrt(16)), "
    "nctx dhead -> qctx dhead)"
)

# Kernel: tile-walk only the lower triangle, online softmax, bias mask
# on the diagonal tile. Per-qctx-tile structural proof.
for iqctx in range(0, qctx.size, qctx_tile_size):
    running_max, running_l, o = -jnp.inf, 0, 0
    for ictx in range(0, iqctx + qctx_tile_size, nctx_tile_size):
        q_tile = Q.slice(qctx, iqctx, iqctx + qctx_tile_size)
        k_tile = K.slice(nctx, ictx, ictx + nctx_tile_size)
        qk = tjax.einsum(q_tile, k_tile, "qctx dhead, nctx dhead -> nctx qctx") / jnp.sqrt(dhead.size)
        # Bias-form mask: 0 inside causal region, -inf outside.
        qk = qk + tjax.mask(qk.type.st, "nctx <= qctx", 0.0, -jnp.inf)
        # ... online-softmax accumulator over qk, V ...
    o.assert_equivalent(SPEC, nctx[:(iqctx + qctx_tile_size)])  # proof
```

The verifier proves the kernel's online softmax, with `-inf` bias on the diagonal tile and skipped upper-triangle tiles, normalizes to the explicit causal-attention spec.

## Installation
Install via pip or uv
```bash
pip install stile-verifier
uv pip install stile-verifier
```

Then import the typed version of your backend!
```python
import stile.jax as tjax # Typed Jax!
import stile.jax.pallas as tpl # Typed Pallas!
import stile.torch as ttorch # Typed Torch!
import stile.triton as ttl # Typed Triton!
```

## How it works

A Stile-typed tensor has two types:
- **`ShapeType`**: which logical dims the tensor has, and how each is sliced.
- **`ExprType`**: the AST of operations performed up to this point.

`ShapeType` is enforced eagerly: you can only multiply two tensors with matching dim signatures, you must reduce the entire dim before assigning the result, etc.

`ExprType` is checked at assignment time. The verifier normalizes both your kernel's accumulated `ExprType` and the spec's parsed `ExprType` into a canonical form and compares. The canonical form folds:
- algebraic identities (`exp(0) = 1`, `0 + x = x`, `0 * x = 0`, `exp(a-b) = exp(a)/exp(b)`, `exp(-inf) = 0`, …),
- distribution through tagged tensors (`*` and `+` push through `Cond(D, …)` branches),
- iteration-domain folding (`sum(body * Cond(D, 1, 0))` collapses to a sum over `[0, N) ∩ D`),
- adjacent-tile interval merging on sum and max reductions, including with shared cross-variable predicates,
- max push-through and `-inf` absorption (so bias-form masks converge with multiplicative-form masks),
- post-fold invariant hoisting (the piece that lets the kernel's `exp(max)` rescaling factor cancel between numerator and denominator).

If the two normalize to the same expression, they compute the same function.

## The specification language

A specification is a small expression over named dims. Tensors are written as the sequence of their dims (`Q D` is a tensor of shape `Q × D`). Slices use `D[a:b]`. Einsums use `,` to separate operands and `->` to give the output shape.

```
Q D K                                       # a tensor with three dims
Q D[0:8]                                    # tensor with D sliced to [0, 8)
2 * Q D                                     # scaled tensor
A B + A B                                   # addition (same-shape required)
exp(Q D)                                    # elementwise unary
(Q D, K D -> Q K)                           # einsum: sum over D
sum[K](Q K)                                 # explicit reduction
sum[K where K <= Q](Q K)                    # iteration-restricted sum (mult-mask)
max[K where K <= Q](Q K)                    # iteration-restricted max (bias-mask)
softmax[K](Q K)                             # softmax along K
softmax[K where K <= Q](Q K)                # causal softmax — restricts both num and den
sum[N](Q N where N >= 4) -> Q               # mult-mask `where`-clause inside sum
```

**`[d where P]`** restricts the iteration of a reduction by the affine predicate `P`. Lowers based on the surrounding op:
- `sum[d where P]` uses a multiplicative mask (`P` zero-elsewhere).
- `max[d where P]` and `softmax[d where P]` use a bias mask (`-inf` elsewhere) — masked positions vanish through `max`'s identity and `exp`'s zero, restricting both the numerator's exp and the denominator's sum.

**`body where P`** (outside a dim annotation) is always a multiplicative mask on `body`. Use it for non-reduction sparsity and bias-on-output patterns. Inside a `sum` it folds into the reduce's domain via mask-extraction; outside it stays as a `Cond` tag.

**Predicates** are conjunctions of affine inequalities over dim names: `<=`, `<`, `>=`, `>`, `==`, plus `+`, `-`, and `int * dim`. Cross-dim predicates (`nctx <= qctx`) are first-class — they ride along in the reduce's domain and survive interval merging.

## Kernel primitives

The kernel side mirrors a small slice of JAX/Numpy. Today's main backend is `stile.jax` (`tjax`); a numpy backend exists for prototyping.

```python
from stile import dim
import stile.jax as tjax

# Dims live in a global registry so specs and kernels share names.
qctx, nctx, dhead = dim('qctx', 128), dim('nctx', 512), dim('dhead', 16)

# Wrap concrete arrays with their dim signature.
Q = tjax.random.normal(key, qctx, dhead)        # has ShapeType (qctx, dhead)

# Slice. Result's st remembers it's [iqctx, iqctx+T).
q_tile = Q.slice(qctx, iqctx, iqctx + 32)

# Einsum. Both shapes and the AST track the contraction.
qk = tjax.einsum(Q, K, "qctx dhead, nctx dhead -> qctx nctx")

# Reductions. Either a method or via einsum.
m = qk.max(nctx)
s = qk.sum(nctx)

# Unary functions on TypedJaxArrays.
e = tjax.exp(qk - m.repeat(nctx))

# Multiplicative mask sugar — score * Cond(P, 1, 0).
masked = score.where("nctx <= qctx")

# Tagged-constant tensor — picks 0/1, 0/-inf, etc.
mult_mask = tjax.mask(score.type.st, "nctx <= qctx")              # 1 inside / 0 outside
bias_mask = tjax.mask(score.type.st, "nctx <= qctx", 0.0, -jnp.inf)  # 0 / -inf

# Rolled loops. Concrete bounds unroll; symbolic bounds emit a parametric reduce.
total = tjax.fori_loop(0, n, lambda i, acc: acc + body(i), init_val=0.0)

# Verify against a spec.
result = tjax.TypedResult(SPEC)
result.assign(o)                          # full-coverage type check
o.assert_equivalent(SPEC, nctx[:K])       # per-tile check with a slice override
result.done()                             # tile-coverage check (no gaps/overlaps)
```

## Status

Working today, with full structural verification:
- **Backends**: `stile.jax` (primary), `stile.numpy` (prototype).
- **Verified kernels**: matmul, online softmax, full flash attention, **tile-walking causal flash attention** (online softmax with bias-mask; structurally proven equivalent to a one-line `softmax[k where k<=q]` spec).
- **Spec features**: einsums, slices, reductions (`sum`, `max`, `softmax`), unary (`exp`, `sin`, `cos`, `sqrt`), multiplicative `where`-clauses, iteration-restricted `[d where P]` annotations, affine predicates with cross-dim references.
- **Kernel features**: slicing, einsum, all the unary/binary ops, `repeat`, `rearrange`, `fori_loop` (concrete-unroll path; symbolic-loop path with parametric reductions), `mask` intrinsic and `.where(...)` sugar.

In progress / future:
- **TypedPallas**: same type discipline, lowering to Pallas for actual GPU/TPU codegen.
- **TypedTorch**: a Torch backend exists but lags the JAX one.

## Why a type system?

Way back in the 1950s, before high-level languages, programs took big arrays of bytes as input and outputted other arrays of bytes. It was up to the programmer to remember in his head which bytes corresponded
to which semantic piece of the program. The invention of type annotations, structures, etc., was a step change in programming productivity because it gave semantic meanings to specific regions of memory.

The current state of the art for numerical programs is not unlike the 1950s mode of programming. We take big multidimensional arrays of floats, and output big multidimensional arrays of floats. The dimensions
are not semantically enforced to be different, and it's up to the programmer to remember the order of dimensions at all times. Mixing them is all too easy. Then assuming a program actually completed, you 
have another big array of floats, and you have no idea if it's right or not. If it doesn't give you the expected result, debugging a numerical program is a massive time-sink. The solution is therefore to
add guardrails to prevent making stupid mistakes, in other words a type system. 


## Running

```bash
uv run pytest tests/
```

Backend extras:

```bash
uv pip install -e ".[jax]"                 # JAX backend
uv pip install -e ".[torch]"               # Torch backend (lagging)
```
