Kina.ai
Menu

Specification

Phase 1, Section 10: Restricted reverse-mode VJPs for StoryLM-0 loss

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

1. What this section commits to

Section 9 compiles log_softmax and take_along_axis and rejects them in the backward build. StoryLM-0 NLL is the composition

- logits.log_softmax(axis).take_along_axis(targets, axis)

so a training-shaped AD graph needs three VJPs:

  1. neg — already a specified unary; the restricted AD pass had no rule.
  2. log_softmax — standard log-domain softmax cotangent.
  3. take_along_axis — scatter the incoming cotangent back onto the class axis of the values tensor.

This section writes those rules for the existing compiler AD pass (Ad_pass.transform_ad). It does not start Phase 2 AD-as-an-effect.

2. What this section does not commit to

3. Pass constraints (restated, not new)

The restricted pass is unchanged except for the new VJP cases:

A legal NLL AD program therefore looks like:

let nll = fn(logits: tensor<F32, [1, 4, 8], @host>) -> tensor<F32, [1, 4], @host> {
  let targets = tensor<I64, [1, 4], @host> in
  let log_probs = logits.log_softmax(axis=2) in
  let picked = log_probs.take_along_axis(targets, axis=2) in
  let loss = -picked in
  loss
} in
let logits = tensor<F32, [1, 4, 8], @host> in
handle nll(logits) with reverseMode

4. neg VJP

neg is Section 5 / Phase 0 unary negation. This section only fills the missing backward case.

v = -a
─────────────────────────
g_a = -g_v

Placement and dtype of g_a equal those of a. The backward node is the existing Unary("neg", …) form.

5. log_softmax VJP

Let y = log_softmax(x) along a compile-time axis (Section 9 §3). The incoming cotangent g has the same shape as y and x.

s = exp(y)                         (softmax(x); computed from the primal)
g_x = g - broadcast(sum(g, axis, keepdims=true), shape(x)) * s

sum and broadcast are the existing reduction / broadcast primitives. exp is the existing unary; it is not given its own VJP in this section.

This is the standard identity g - sum(g) * softmax. Using exp(y) avoids recomputing the stable max-shift in the backward graph.

VJP placement follows x. Device log_softmax may therefore produce a device backward graph composed of ops that already have device lowering. No new FFI.

6. take_along_axis VJP

Write σ = α ++ [C] ++ β with |α| = axis, and

out[a ++ b] = values[a ++ [indices[a ++ b]] ++ b]

as in Section 9 §4.2. The incoming cotangent g has shape ρ = α ++ β.

g_values[a ++ [c] ++ b] =
  g[a ++ b]   if indices[a ++ b] = c
  0           otherwise

Each destination point (a, c, b) is written by at most one source (a, b). Duplicate class ids at different (a, b) land at different outputs. A parallel compare/select over the destination shape is therefore enough; this slice does not need atomics or a sequential scatter-add.

The compiler represents that cotangent with the IR-internal constructor in §7. It does not extend 2D gather.

7. ScatterAlongAxis (IR-internal)

7.1 Typing

Not a surface form. The typed IR constructor is

ScatterAlongAxis of int * program * program * tensor_type

in constructor order (axis, updates, indices, dest_ty).

Γ ⊢ updates : Tensor[T, ρ]
Γ ⊢ indices : Tensor[I64, ρ]
0 ≤ axis < |σ|
|σ| ≥ 2
ρ ≡ σ with the dimension at `axis` removed
updates.placement = indices.placement = dest_ty.placement
dest_ty.dtype = T
dest_ty.shape = σ
placement = Host                         (this slice)
∀i. 0 ≤ indices[i] < σ[axis]             (runtime-checkable bound)
─────────────────────────────────────────
Γ ⊢ ScatterAlongAxis(axis, updates, indices, dest_ty) : Tensor[T, σ]

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

The StoryLM-0 instance:

updates : Tensor[F32, [B, T], @host]
indices : Tensor[I64, [B, T], @host]
axis    = 2
dest_ty : Tensor[F32, [B, T, V], @host]
result  : Tensor[F32, [B, T, V], @host]

The type checker does not accept a surface scatter_along_axis. The AD pass constructs well-typed nodes. infer_type / typeof_program read dest_ty. Device operands are rejected at emit, same stance as take_along_axis.

7.2 Operational meaning

out[a ++ [c] ++ b] =
  updates[a ++ b]   if indices[a ++ b] = c
  0                 otherwise

Out-of-range indices trap on the host path (cf.assert), same policy as checked take_along_axis.

7.3 Host lowering

  1. A bounds linalg.generic over indices (compare against 0 and C), assertion message scatter_along_axis index out of bounds.
  2. A linalg.generic over the destination shape σ. Affine maps project away axis for updates and indices. The body compares linalg.index axis (cast to i64) with the index element and arith.selects the update versus a typed zero.

Device placement fails emit: scatter_along_axis: checked device lowering is not available.

The forward AD walk rewrites children of ScatterAlongAxis if such a node appears. Building a VJP for ScatterAlongAxis 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. Tests

9. Open issues

  1. reverseMode result ABI. Specified in Section 12. The language type is the flat tuple (primal, g₀, …); @main / fn use MLIR multi-result returns. Tuple projection (loss → VJP → SGD in one program) remains open there.
  2. Device take_along_axis / ScatterAlongAxis. Still follow-on.
  3. ScatterAlongAxis VJP (would be a take_along_axis of the dest cotangent). Not needed while scatter exists only in the backward graph.
  4. User-facing exp VJP. Not in this section.
  5. Mean NLL. Optional reduce of the per-site loss; existing sum VJP already covers a subsequent reduction.