Phase 1, Section 14: Restricted reverse-mode VJPs for the StoryLM-0 decoder graph

Compiler transform_ad only. Not Phase 2 handler AD. Not a Phase 1 exit-criteria replacement. These rules extend Sections 5 and 10; they do not add a user-facing scatter primitive.

1. What this section commits to

Section 10 differentiates neg, log_softmax, and take_along_axis. The StoryLM-0 forward graph in storylm0_forward.kina still fails the restricted AD pass because three already-specified forward operations have no backward rule:

1. reshape — Phase 0 §2 volume-preserving reshape.
2. Variadic-batch matmul — Section 5 §3. Rank 2 and rank 3 already have VJPs; rank 4 is the attention case [B, H, M, K] @ [B, H, K, N].
3. 2D gather — Section 5 §4 embedding lookup.

This section writes those VJPs for Ad_pass.transform_ad. It also records one pass clarification required to inject token ids through a host .so: discrete I64 parameters carry no cotangent and are omitted from the reverseMode result tuple.

The StoryLM-0 numeric instance is one program that runs the decoder graph, names the embedding cotangent, and applies Section 11 SGD.

2. What this section does not commit to

• Phase 2 handler AD, tape AD, or differentiating an arbitrary program.
• A user-facing scatter, scatter_add, one_hot, or arange.
• 3D-values gather (Section 5 open issue 3). gather stays Tensor[T, [V, D]].
• Device gather / ScatterAdd AD (host only, same stance as Section 10 ScatterAlongAxis).
• Differentiating through a destructuring let (Section 13).
• AdamW, argmax, seeded sampling, or KV-cache semantics.
• Training every weight in one SGD step. The instance updates token_embedding only; other floating parameters still receive cotangents and may be bound unused.
• Changing the seed cotangent (ones of the primal result).

3. Pass constraints (restated, with one amendment)

Unchanged from Section 10:

• The differentiated function body is a let chain that ends with a Var.
• Each differentiated node’s children are Vars. Method chains such as x.reshape(σ).transpose(π) must be split into two lets.
• The seed cotangent is ones of the primal result shape.
• A let that never receives a cotangent is skipped, not an error.

Amendment (discrete parameters). Section 10 required every reverseMode function parameter to receive a cotangent, and therefore forced I64 indices into inner lets. Inner-let tensors cannot be injected through the host C wrapper (payload { … } is discarded). This section amends the rule:

• A parameter whose dtype is not a floating tensor (F16, BF16, F32, F8E4M3, F8E5M2) is discrete. It is not entered in the adjoint table and does not appear in the result tuple.
• Floating parameters still each receive a cotangent, in parameter order, after the primal.
• The result type is therefore (R, g₀, …, gₖ₋₁) where k is the number of floating parameters. If k = 0, the result is R (Section 12 zero-parameter collapse).

A legal decoder AD program may therefore take token_ids : Tensor[I64, σ] as a parameter. The host .so can write those ids. They still have no adjoint.

4. reshape VJP

Phase 0 §2 reshape is a volume-preserving, element-order rearrangement. Write y = reshape(x, σ') with vol(shape(x)) = vol(σ') already discharged. The incoming cotangent g has shape σ'.

g_x = reshape(g, shape(x))

Placement and dtype of g_x equal those of x. The backward node is the existing Reshape. No new primitive.

This is the unique linear inverse of a volume-preserving reshape: the same elements, original layout.

5. Variadic-batch matmul VJP

Section 5 §3 types

a : Tensor[T, σ ++ [M, K]]
b : Tensor[T, σ ++ [K, N]]
a @ b : Tensor[T, σ ++ [M, N]]

for any leading σ (no implicit broadcast). Rank 2 is σ = []; rank 3 is |σ| = 1; the StoryLM-0 attention matmuls are |σ| = 2.

The VJP is the same contraction already implemented for rank 2 and 3, written with a last-two-dimension transpose:

aᵀ = transpose_last_two(a)     : σ ++ [K, M]
bᵀ = transpose_last_two(b)     : σ ++ [N, K]
g_b = aᵀ @ g                   : σ ++ [K, N]
g_a = g @ bᵀ                   : σ ++ [M, K]

transpose_last_two is the existing Transpose whose permutation is the identity on σ and swaps the last two axes:

perm = [0, 1, …, r−3, r−1, r−2]     (r = |σ| + 2 ≥ 2)

Rank 2: perm = [1, 0]. Rank 3: perm = [0, 2, 1]. Rank 4: perm = [0, 1, 3, 2]. Rank < 2 is not a matmul.

The backward nodes are existing Transpose and Matmul. Host rank-4 matmul already lowers through tensor.collapse_shape / linalg.batch_matmul / tensor.expand_shape (Section 5 §3.2). The backward graph reuses that path.

This section does not add mixed-precision matmul (Section 5 open issue 4).

6. gather VJP

Section 5 §4:

values  : Tensor[T, [V, D]]
indices : Tensor[I64, σ]
y       : Tensor[T, σ ++ [D]]
y[i ++ [d]] = values[indices[i], d]

The incoming cotangent g has shape σ ++ [D].

• indices is discrete: no cotangent (parameter or inner let).
• values cotangent has shape [V, D]:

g_values[v, d] = Σ_{i : indices[i] = v} g[i ++ [d]]

Several positions in σ may share a vocab id. Contributions add. This is not Section 10 ScatterAlongAxis:

The compiler represents g_values with the IR-internal constructor in §7. It does not extend 2D gather and does not reuse ScatterAlongAxis.

7. ScatterAdd (IR-internal)

7.1 Typing

Not a surface form. The typed IR constructor is

ScatterAdd of program * program * tensor_type

in constructor order (updates, indices, dest_ty).

Γ ⊢ updates : Tensor[T, σ ++ [D]]
Γ ⊢ indices : Tensor[I64, σ]
|σ| ≥ 1
dest_ty.dtype = T
dest_ty.shape = [V, D]
updates.placement = indices.placement = dest_ty.placement
placement = Host                         (this slice)
∀i. 0 ≤ indices[i] < V                   (runtime-checkable bound)
─────────────────────────────────────────
Γ ⊢ ScatterAdd(updates, indices, dest_ty) : Tensor[T, [V, D]]

dest_ty is the result type (the primal values type of the corresponding gather). Structural shape equality, same as Section 9.

The StoryLM-0 instance:

updates : Tensor[F32, [1, 4, 4], @host]
indices : Tensor[I64, [1, 4], @host]
dest_ty : Tensor[F32, [8, 4], @host]
result  : Tensor[F32, [8, 4], @host]

The type checker does not accept a surface scatter_add. The AD pass constructs well-typed nodes. infer_type / typeof_program read dest_ty. Device operands are rejected at emit.

7.2 Operational meaning

out[v, d] = Σ_{i : indices[i] = v} updates[i ++ [d]]

Rows that no index selects are zero. Out-of-range indices trap on the host path (cf.assert), same policy as checked gather.

7.3 Host lowering

1. A bounds linalg.generic over indices (compare against 0 and V), assertion message scatter_add index out of bounds.
2. Flatten σ to N = Π σ when |σ| ≥ 2 (tensor.collapse_shape on indices → [N] and on updates → [N, D]). Rank-1 σ is already flat. This slice requires static dims for the flatten.
3. Zero-initialize a dest tensor of type dest_ty.
4. A linalg.generic over iterators (v, d, n) with types parallel, parallel, reduction:

out[v, d] +=  (indices_flat[n] == v) ? updates_flat[n, d] : 0

No atomics. Duplicates accumulate because n is a reduction iterator. Device placement fails emit: scatter_add: checked device lowering is not available.

The forward AD walk rewrites children of ScatterAdd if such a node appears. Building a VJP for ScatterAdd is unsupported in this slice (it is only inserted in the backward graph).

7.4 Surface

None. No .kina method, no S-expression keyword.

8. StoryLM-0 numeric instance

The decoder body is the storylm0_forward.kina graph with every operator in its own let (the existing method-chain form is not an AD spine). token_ids is an I64 parameter (no cotangent). 2D weights add a unit batch dim with reshape, not rank-changing broadcast (that VJP is still unspecified). The wrapper applies Section 11 SGD to the embedding table:

let (logits, g_emb, g_mask, g_scale, g_wq, g_wk, g_wv, g_wo, g_lm) =
  handle storylm0_fwd(token_ids, emb, mask, scale, wq, wk, wv, wo, lm_head)
  with reverseMode in
emb - lr.broadcast([8, 4]) * g_emb

Shapes match the forward example: emb : [8, 4], lr : [1, 1], logits [1, 4, 8]. The seed is ones on the logits, so g_emb is ∂(sum logits) / ∂emb. Unused cotangents (g_mask, …) are bound and discarded.

just validate-decoder-ad compiles isolated gather / reshape / rank-4 matmul wrappers and this train step to .so files and checks them against NumPy (PyTorch when present).

9. Tests

• ad_reshape.kina — reshape VJP compiles; wrapper returns reshape(ones, original_shape).
• ad_matmul4.kina — rank-4 matmul VJP compiles; wrapper returns g_q = ones @ kᵀ for StoryLM-0 attention shapes [1, 2, 4, 2] and [1, 2, 2, 4].
• ad_gather.kina — 2D gather VJP compiles; emitted graph contains ScatterAdd (assert scatter_add index out of bounds and a reduction-iterator add). An I64 index parameter does not appear in the reverseMode tuple.
• storylm0_ad_step.kina — §8 program; golden accepted by mlir-opt; the step fn result is tensor<8x4xf32>.
• Structural IR: after transform_ad, gather inserts ScatterAdd and does not insert ScatterAlongAxis.
• scripts/validate_decoder_ad.py: isolated VJPs and the embedding SGD step match a trusted reference.

10. Open issues

1. Device ScatterAdd / gather AD. Follow-on, same as device take_along_axis.
2. ScatterAdd VJP (would be a gather of the dest cotangent). Not needed while scatter-add exists only in the backward graph.
3. SGD on every decoder weight. Specified as composition in Section 15.
4. NLL ∘ decoder in one program. Specified as composition in Section 15.
5. AdamW, argmax, sampling, KV-cache. Still unspecified.
6. Symbolic-shape flatten in ScatterAdd host lowering.
7. Rank-changing broadcast VJP. Same-rank broadcast (expand 1s) already has a VJP. Tensor[T, [M, K]] → [1, M, K] is not that rule; this slice uses volume-preserving reshape instead.