ADR: StoryLM-0 decoder-graph VJPs

Context

Section 10 fills VJPs for the NLL composition. The StoryLM-0 forward graph still uses gather, reshape, and rank-4 matmul, which the restricted compiler AD pass rejects. Rank-2/3 matmul already has a last-two-dimension transpose VJP. Section 5 already types variadic-batch matmul and 2D gather. Phase 0 already types volume-preserving reshape.

take_along_axis uses IR-internal ScatterAlongAxis (compare/select, no add). Embedding gather can map many tokens to one vocab row, so its cotangent must add.

Host .so checks need injectable token ids. Section 10 forced I64 indices into inner lets, whose payloads cannot be written from C.

Decision

Specify and implement the VJPs in 14-decoder-ad.md:

• reshape: g_x = reshape(g, shape(x)).
• Variadic-batch matmul (rank ≥ 2): g_b = aᵀ @ g, g_a = g @ bᵀ with last-two-dimension transpose. Rank 4 is the attention instance.
• gather: no index adjoint; values cotangent is IR-internal ScatterAdd (reduction-iterator add over flattened σ). Not ScatterAlongAxis.
• Discrete (I64 and other non-floating) reverseMode parameters are omitted from the result tuple so token ids can be function arguments.

The StoryLM-0 numeric instance names g_emb with Section 13 projection and applies Section 11 SGD to the embedding table. Other floating cotangents may be bound unused.

Consequences

• The typed IR grows one constructor that the parser never produces.
• Host MLIR goldens for ad_gather, ad_reshape, ad_matmul4, and storylm0_ad_step land with the pass.
• Device scatter-add and Phase 2 handler AD remain follow-on work.
• Multi-weight SGD and NLL∘decoder are specified as composition in Section 15.