Phase 0 Casebook: Real-World Scenarios

A different kind of companion: not a long survey, but a hands-on casebook.

Use this when you want to answer: "If I wrote real model code, what exactly would I do in Phase 0?"

If this file and phase-0/ ever disagree, phase-0/ wins.

Quick start

Pick one case below and follow this loop:

1. Read the intent in plain English.
2. Check the shape ledger table (input -> intermediate -> output).
3. Read the Phase 0 implementation sketch.
4. Run the checkpoint questions at the end.

Case A: RMSNorm in production-style code

Intent

Normalize each token vector and rescale it with learned weights.

Shape ledger

Phase 0 implementation sketch

let rmsnorm [B, S, D]
    (x : Tensor[F32, [B, S, D]])
    (weight : Tensor[F32, [D]])
    (eps : Tensor[F32, []])
    : Tensor[F32, [B, S, D]] =
  let x_sq         = x * x in
  let mean_sq      = reduce_keep[mean, axis=2] x_sq in
  let eps_bcast    = broadcast eps [B, S, 1] in
  let denom_inner  = mean_sq + eps_bcast in
  let rstd         = rsqrt denom_inner in
  let rstd_bcast   = broadcast rstd [B, S, D] in
  let normed       = x * rstd_bcast in
  let weight_bcast = broadcast weight [B, S, D] in
  normed * weight_bcast

Why this case matters

• Shows explicit broadcast (core Phase 0 rule).
• Shows reduce_keep pattern you will reuse everywhere.
• No fancy typing tricks needed; this is the "happy path".

Checkpoint questions

• Why is weight not directly multiplied without broadcast?
• What would break if reduce_keep were plain reduce here?

Case B: Attention score path (where complexity appears)

Intent

Compute softmax((Q K^T) * scale) V for shapes [B, H, S, D_h].

Reality in Phase 0

You usually need these explicit steps:

1. transpose keys.
2. Flatten [B, H] to one batch-like dim for current matmul form.
3. matmul to get scores.
4. Explicit softmax decomposition:
   • row max
   • subtract max (stability)
   • exp
   • row sum
   • divide
5. matmul with V.
6. Unflatten back to [B, H, S, D_h].

Practical takeaway

This is not a beginner failure; it is an expected Phase 0 limitation. It directly motivates:

• variadic-batch matmul,
• softmax as primitive or robust sugar,
• better surface ergonomics in Phase 7.

Case C: What user code still looks like at block level

Even with verbose primitives, composition becomes readable:

let transformer_block_partial [B, S, D, H] ... : Tensor[F32, [B, S, D]] =
  let h_norm = rmsnorm[B, S, D] x norm_w eps in
  let h_ffn  = swiglu_ffn[B, S, D, H] h_norm w_gate w_up w_down in
  x + h_ffn

Meaning: keep complexity inside library functions; application code stays small.

Debug checklist for any new case

When a case is hard to type-check, check these in order:

1. Did dtype change? If yes, you likely need cast.
2. Did shape rank/order change? If yes, you likely need explicit reshape/transpose.
3. Are you relying on implicit broadcast? Add explicit broadcast.
4. Is the reduce axis a literal? Phase 0 expects compile-time literal axis.
5. Did you use shape arithmetic beyond Presburger limits? Use documented escape hatch.

Next practice (30 minutes)

• Recreate Case A from memory and verify each shape line-by-line.
• Write a mini function residual_add(x, y) and prove to yourself why both inputs must have identical shape in core IR.
• Take one PyTorch helper you already know and map it to explicit Phase 0 operations.

Related docs

• companions/phase-0-learning-path.md
• companions/examples-walkthrough.md
• companions/real-world-examples.md
• phase-0/02-core-operations-and-typing.md