Learn

Phase 0 Companion: Example Programs Walkthrough

Companion to the Phase 0 specification. Not part of the formal spec; an illustration of what programs look like under the rules Phase 0 commits to.

1. Framing

Every example below appears in two forms:

Surface (illustrative). What a programmer would actually write. The syntax is OCaml-flavored and speculative — Phase 7 commits to a final surface; this is just plausible-looking syntax to show what the type system feels like.
Core IR. What Phase 0 actually commits to. S-expression form; what the OCaml type-checker consumes after parsing and desugaring.

The point of the walkthrough: the formal rules in Sections 1–4 are precise but abstract. Programs make them tangible — including the costs of the design decisions we made (no implicit broadcasting, restricted Presburger refinements, the controlled unsoundness of reshape).

If a program below feels unreasonably verbose, that's a real signal about a design tradeoff. A few of them do, and §10 returns to that.

2. Example 1: Element-wise addition

The simplest non-trivial program. Two same-shape tensors, add them.

Surface

let add2d [N, M]
    (x : Tensor[F32, [N, M]])
    (y : Tensor[F32, [N, M]])
    : Tensor[F32, [N, M]] =
  x + y

Core IR

(fn add2d
  (forall (N : Dim) (M : Dim))
  ((x (Tensor F32 [N M] []))
   (y (Tensor F32 [N M] [])))
  (Tensor F32 [N M] [])
  (binop add x y))

What the type-checker does

Parse → AST → elaborate to core IR.
At the function body: synthesize (binop add x y) against the elementwise-binary rule (Section 2 §6).
The rule requires x : Tensor[F32, s, R₁] and y : Tensor[F32, s, R₂] with the same s. Both are [N, M]. Match.
Output type: Tensor[F32, [N, M], [] ∧ []] = Tensor[F32, [N, M], []].
Subsume against the annotated return type: trivially holds.

No SMT calls. No counterexamples. Compiles to a single MLIR linalg.add lowering to LLVM. Done.

3. Example 2: Linear layer (matmul + bias)

A fully-connected layer applied to a batched 3D input. This is where the no-implicit-broadcasting decision starts costing.

Surface (illustrative; assumes Phase 7 sugars matmul as `@`)

let linear [B, S, D_in, D_out]
    (x : Tensor[F32, [B, S, D_in]])
    (w : Tensor[F32, [D_in, D_out]])
    (b : Tensor[F32, [D_out]])
    : Tensor[F32, [B, S, D_out]] =
  let w_batched = broadcast w [B, D_in, D_out] in
  let h = x @ w_batched in
  let b_expanded = broadcast b [B, S, D_out] in
  h + b_expanded

Core IR

(fn linear
  (forall (B : Dim) (S : Dim) (D_in : Dim) (D_out : Dim))
  ((x (Tensor F32 [B S D_in] []))
   (w (Tensor F32 [D_in D_out] []))
   (b (Tensor F32 [D_out] [])))
  (Tensor F32 [B S D_out] [])
  (let
    ((w_batched (broadcast w [B D_in D_out]))
     (h         (matmul x w_batched))
     (b_exp     (broadcast b [B S D_out])))
    (binop add h b_exp)))

What's verbose here

Every shape change is explicit. The user wrote:

One broadcast to give w a batch dim.
One matmul.
Another broadcast to expand b from [D_out] to [B, S, D_out].
One add.

In PyTorch this would be x @ w + b — three of those four operations are implicit. Phase 7's surface can recover them, but at the IR level they're four real operations with four real lowerings.

The cost is real. The benefit is that lowering sees them all and can fuse aggressively (the broadcast-then-matmul fuses to a non-broadcasting batched-matmul; the broadcast-then-add fuses to a single bias-add kernel). Trading user friction at the IR level for unambiguous lowering is the deal Phase 0 made.

4. Example 3: MLP forward pass

Two linear layers with a ReLU between them. Small but realistic.

Surface

let mlp_forward [B, S, D_in, D_hidden, D_out]
    (x  : Tensor[F32, [B, S, D_in]])
    (w1 : Tensor[F32, [D_in, D_hidden]])
    (b1 : Tensor[F32, [D_hidden]])
    (w2 : Tensor[F32, [D_hidden, D_out]])
    (b2 : Tensor[F32, [D_out]])
    : Tensor[F32, [B, S, D_out]] =
  let h1 = linear x w1 b1 in
  let h1_act = relu h1 in
  let h2 = linear h1_act w2 b2 in
  h2

What's interesting

Once linear is defined (Example 2), the MLP composes cleanly. All the shape arithmetic happens at the linear-layer boundary; the MLP itself reads almost like Python.

The type-checker, at each call to linear, instantiates the polymorphic shape variables from the call site:

First call: B, S, D_in, D_hidden — bound by passing x : [B, S, D_in] and w1 : [D_in, D_hidden].
Second call: B, S, D_hidden, D_out — bound by passing h1_act : [B, S, D_hidden] and w2 : [D_hidden, D_out].

No SMT calls (all shape relationships are syntactic). No counterexamples. The whole MLP type-checks in microseconds.

5. Example 4: A piece of multi-head attention

Where the verbosity costs really show up. Not the full attention — softmax is a composite (exp + sum + div + broadcast) and we don't yet have a way to abstract it cleanly. Just the score computation Q @ K^T.

Surface

let attention_scores [B, H, S_q, S_k, D_h]
    (q : Tensor[F32, [B, H, S_q, D_h]])
    (k : Tensor[F32, [B, H, S_k, D_h]])
    : Tensor[F32, [B, H, S_q, S_k]] =
  let k_t = transpose k [0, 1, 3, 2] in
  // K^T : [B, H, D_h, S_k]
  let dim BH = B * H in
  let q_flat   = reshape q   [BH, S_q, D_h] in
  let k_t_flat = reshape k_t [BH, D_h, S_k] in
  let scores_flat = q_flat @ k_t_flat in
  // scores_flat : [BH, S_q, S_k]
  reshape scores_flat [B, H, S_q, S_k]

What's happening

Three things conspire to make this verbose:

Matmul is batched with exactly one batch dim (Section 2 §7). Q has two leading batch dims (B, H); K has the same. We have to flatten them into one before calling matmul, then unflatten after.
The dim BH = B * H line introduces a fresh shape variable with the controlled-escape-hatch from Section 1 §3. Variable-times-variable is outside Presburger, so the equality BH = B * H is asserted axiomatically — the type-checker accepts it without proving it. The runtime catches a mistake (wrong volume in a reshape) only at execution time. Documented unsoundness.
Two reshapes — one to flatten, one to unflatten. Each carries its own volume axiom. Phase 0 doesn't prove that the unflatten is the inverse of the flatten; it just trusts the user.

Core IR (abbreviated — full form is similar)

(fn attention_scores
  (forall (B : Dim) (H : Dim) (S_q : Dim) (S_k : Dim) (D_h : Dim))
  ((q (Tensor F32 [B H S_q D_h] []))
   (k (Tensor F32 [B H S_k D_h] [])))
  (Tensor F32 [B H S_q S_k] [])
  (let
    ((k_t        (transpose k [0 1 3 2]))
     (BH         (fresh-dim))                  ; introduce variable
     (axiom      (assert-eq BH (mul B H)))     ; axiomatically
     (q_flat     (reshape q   [BH S_q D_h]))
     (k_t_flat   (reshape k_t [BH D_h S_k]))
     (scores_flat (matmul q_flat k_t_flat))
     (scores     (reshape scores_flat [B H S_q S_k])))
    scores))

Reaction

This is the most verbose example so far, and it's only one piece of attention. The full attention layer is roughly four times this, plus softmax, plus value matmul, plus output projection.

Two honest takes:

Phase 7 surface syntax recovers most of the verbosity. Pythonic syntax with operator-overloaded @, automatic batch flattening for matmul, and named-tensor support could turn this into:
```
scores = (q @ k.transpose(-2, -1))   -- maybe 1 line
```
That's a Phase 7 elaboration concern; the core IR stays explicit.
Some of the verbosity is fundamental. The BH = B * H flattening is a real consequence of Phase 0's choice to keep matmul batched with one batch dim and to keep multiplication out of the refinement language. Either of those choices could be revisited if attention-class code is the dominant case (and it is, for ML). It might be worth reconsidering both for Phase 1.

This is exactly the kind of tradeoff the walkthrough is supposed to surface. The spec says the right things; the program shows what the spec costs.

6. Example 5: Type error — shape mismatch

A program that should be rejected. Demonstrates what failures look like.

Surface (intentionally wrong)

let bad_add [B, S, D]
    (x : Tensor[F32, [B, S, D]])
    (y : Tensor[F32, [B, S, D + 1]])
    : Tensor[F32, [B, S, D]] =
  x + y

Compiler output (illustrative)

Error at bad_add:5,7
  Type error in elementwise binary operation `add`.
  Expected: shapes equal at every position.
  At position 2:
    Left:  D
    Right: D + 1
  
  SMT counterexample: D = 1
    Under this assignment, left shape position 2 = 1, right = 2.
  
  Hint: insert an explicit broadcast or reshape to align shapes.

Build failed (1 error).

What happened internally

The type-checker reaches (binop add x y).
The elementwise-binary rule requires x and y to have exactly matching shapes.
Shape unification: position 0 (B = B) is syntactically equal — pass. Position 1 (S = S) — pass. Position 2 (D = D + 1) — call SMT.
SMT: assert Γ ∧ ¬(D = D + 1). Always satisfiable (e.g., D = 1 gives D = 1, D + 1 = 2, and ¬(1 = 2) is true).
Z3 returns sat with a counterexample. Type-checker reports the failure with the counterexample.

The counterexample is what makes refinement-typed errors usable. "Shape mismatch" alone is unhelpful; "shape mismatch, here's an assignment showing why" is debuggable.

7. Example 6: Refinement violation across function calls

A more interesting error: the call site violates a requires clause.

Surface

let split_heads [B, S, D, H]
    requires (D mod H = 0)
    (x : Tensor[F32, [B, S, D]])
    : Tensor[F32, [B, H, S, D / H]] =
  let dim D_per_head = D / H in
  let reshaped = reshape x [B, S, H, D_per_head] in
  transpose reshaped [0, 2, 1, 3]

let bad_caller [B, S]
    (x : Tensor[F32, [B, S, 100]])
    : Tensor[F32, [B, 7, S, 14]] =
  split_heads[B, S, 100, 7] x

Compiler output

Error at bad_caller:11,3
  Cannot satisfy precondition of `split_heads`.
  Required: D mod H = 0.
  At call site: D = 100, H = 7.
  
  SMT counterexample: 100 mod 7 = 2 ≠ 0.
  
  Hint: choose H such that H divides 100,
        or reshape x to a head-divisible hidden dim.

Build failed (1 error).

What happened internally

Type-checker reaches the call to split_heads.
Substitutes B ↦ B, S ↦ S, D ↦ 100, H ↦ 7 into the function's signature.
Discharges the requires clause: Γ ⊨ 100 mod 7 = 0.
SMT: assert ¬(100 mod 7 = 0). Trivially true (100 mod 7 = 2). Z3 returns sat.
The discharge fails. Type-checker reports the violation.

This is the case where requires clauses earn their keep. The function split_heads is guaranteed to be called only with D-divisible-by-H tensors; downstream code can rely on the fact that D / H is exact (no rounding) without further checks. The type system enforces it at the boundary, not inside the function body.

8. Example 7: A program that uses refinement subtyping

Shows where Section 3's subtyping does real work — flowing a more-constrained tensor through a less-constrained interface.

Surface

// Function expecting any non-empty 1D tensor.
let normalize [N]
    requires (N >= 1)
    (x : Tensor[F32, [N]])
    : Tensor[F32, [N]] =
  // ... some normalization computation ...
  x  // placeholder

// Function that produces a tensor known to have at least 16 elements.
let make_aligned_tensor [N]
    requires (N >= 16)
    : Tensor[F32, [N]] =
  // ... some construction ...
  ???  // placeholder

// Calling normalize with an aligned tensor.
let pipeline [N]
    requires (N >= 16)
    : Tensor[F32, [N]] =
  let t = make_aligned_tensor[N] in
  normalize[N] t   // <-- here

What the type-checker does at the call site `normalize[N] t`

t has type Tensor[F32, [N], [N >= 16]].
normalize expects Tensor[F32, [N], [N >= 1]].
Subtyping check: is Tensor[F32, [N], [N >= 16]] <: Tensor[F32, [N], [N >= 1]]?
Decomposes by Section 3 §3.3:
- Dtype equal: ✓
- Shape equal: ✓
- Refinement implication: Γ; (N >= 16) ⊨ (N >= 1)?
SMT: assert Γ ∧ (N >= 16) ∧ ¬(N >= 1) = (N >= 16) ∧ (N < 1). Unsatisfiable.
Implication holds. Subtyping accepted. Subsumption applies. Call type-checks.

No coercion is inserted. t flows through as itself. The "refinement got weaker" was a free upcast.

9. What's missing for a "real" ML program

Reasonable question: can you write a transformer in Phase 0?

Forward-pass-wise, mostly yes — with effort. A transformer block is matmul, layer norm, attention, MLP, and residual connections. Layer norm is reduce + broadcast + elementwise operations; doable. Attention's score computation is Example 4; the softmax is exp + reduce sum + broadcast + binop div, all in Phase 0's operation set. Multi-head attention plus an MLP plus residuals is ~150 lines of core IR.

Backward-pass, no. Reverse-mode AD is Phase 2 (effects). Phase 0 has no AD.

Training loop, no. Mutable parameters need state, an effect. Phase 2.

Distribution / multi-GPU, no. Sharding is Phase 3.

Custom kernels (FlashAttention etc.), no. Hardware-aware tile programming is Phase 5.

Sparsity (MoE, structured sparse), no. Phase 4.

A working forward pass that you could run on a single CPU and verify against PyTorch, yes. That's the Phase 1 deliverable.

So Phase 0 + Phase 1 gives you "type-checked forward passes that compile and run." That's enough to validate the framework and to start writing real benchmarks. Phases 2–6 add what's needed to compete with production ML stacks.

10. Reactions to surface in light of these examples

Three things this walkthrough surfaces that the spec doesn't:

10.1 The matmul batched-with-one-batch-dim choice may be wrong

Example 4's flattening dance is awkward and would happen at every attention call site. Real production code has tensors with 4+ leading dimensions all the time ([batch, head, sequence, hidden], [batch, head, kv_head, sequence, hidden] for grouped-query attention).

A revised matmul that allows any number of leading batch dims (via a ... rest pattern in the type) would be much cleaner. The cost is a slightly more complex typing rule. This is worth reconsidering in Phase 1.

10.2 The variable-product escape hatch will be common

BH = B * H will appear in every attention layer, every grouped-query attention, every MoE routing computation. The escape hatch works but reads as visibly-unsound boilerplate.

Two possible improvements:

A built-in flatten / unflatten operation pair that handles the volume axiom internally, with a typing rule that's specifically designed for the inverse-of-each-other case. This would let the user write flatten(x, axes=[0,1]) and unflatten(y, axes=[0,1], sizes=[B, H]) instead of two reshapes plus an asserted equality.
Extending the refinement language to handle products of bound dimensions in restricted forms. Decidable nonlinear arithmetic isn't open in general, but the patterns we need (products of bound shape variables) are simpler than general nonlinear. Worth investigating.

Phase 0 commits to the escape hatch. Phase 1 should reconsider whether to add flatten/unflatten as primitives.

10.3 The verbosity case for Phase 7 is even stronger than the spec made it

Reading Example 4, Phase 7's surface syntax isn't a "nice-to-have" — it's a "this language is unusable without it." The type system is fine; the IR is fine; the human ergonomics depend entirely on the surface that hasn't been designed yet.

This shifts the project's risk model. Phase 0 gives you the foundation. Phase 1 gives you a working compiler. Phase 7 (surface syntax) is what turns it into something usable — and the gap between Phase 1 and Phase 7 is real years of work even after Phase 1 succeeds.

For a research artifact (papers, prototypes), Phase 1 is enough. For a language, Phase 7 is the one that actually matters.

11. Summary

Seven examples; six well-typed, two with errors. Forward-pass code is expressible; verbosity is real and partly recoverable via Phase 7 surface, partly indicating real design tradeoffs to revisit in Phase 1.

The biggest revelations from the walkthrough:

The matmul-with-one-batch-dim rule needs to be reconsidered.
A flatten/unflatten primitive pair is probably worth adding to the core operation set.
Phase 7 surface syntax is more critical to the project's success than Phase 0 made it sound.

None of these break Phase 0. All of them inform Phase 1.

End of walkthrough.

Phase 0 Companion: Example Programs Walkthrough

1. Framing

2. Example 1: Element-wise addition

Surface

Core IR

What the type-checker does

3. Example 2: Linear layer (matmul + bias)

Surface (illustrative; assumes Phase 7 sugars matmul as @)

Core IR

What's verbose here

4. Example 3: MLP forward pass

Surface

What's interesting

5. Example 4: A piece of multi-head attention

Surface

What's happening

Core IR (abbreviated — full form is similar)

Reaction

6. Example 5: Type error — shape mismatch

Surface (intentionally wrong)

Compiler output (illustrative)

What happened internally

7. Example 6: Refinement violation across function calls

Surface

Compiler output

What happened internally

8. Example 7: A program that uses refinement subtyping

Surface

What the type-checker does at the call site normalize[N] t

9. What's missing for a "real" ML program

10. Reactions to surface in light of these examples

10.1 The matmul batched-with-one-batch-dim choice may be wrong

10.2 The variable-product escape hatch will be common

10.3 The verbosity case for Phase 7 is even stronger than the spec made it

11. Summary

Surface (illustrative; assumes Phase 7 sugars matmul as `@`)

What the type-checker does at the call site `normalize[N] t`