Phase 1, Section 9: `log_softmax` and `take_along_axis`

StoryLM-0 loss slice. Not a Phase 1 exit-criteria replacement. These rules extend Section 5; they do not change 2D gather or softmax.

1. What this section commits to

StoryLM-0 already compiles a causal forward graph to [B, T, V] logits (storylm0_forward.kina). Training loss needs two things the Section 5 operation set does not provide:

1. A stable log-domain softmax on a compile-time axis.
2. Per-position class selection from a rank-n tensor along that axis, so targets of shape σ without the class dimension pick one logit per site.

This section adds those two primitives. Negative log-likelihood is not a third primitive. It is the composition

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

Tokenization, batching, and SGD are out of scope. Reverse-mode VJPs for these ops and for NLL’s neg are specified in Section 10.

2. Why not extend gather or softmax

Phase 1 gather (Section 5 §4) is Tensor[T, [V, D]] × Tensor[I64, σ] → Tensor[T, σ ++ [D]]. That is embedding lookup. Selecting class targets[b, t] from logits[b, t, :] would require 3D-values gather, which Section 5 §9.3 explicitly deferred.

softmax already exists. log_softmax is a separate primitive so the compiler emits the stable x - max - log(sum(exp(x - max))) form instead of log(softmax(x)), which underflows.

3. log_softmax

3.1 Typing rule

Same premises as Section 5 §5.1 softmax:

Γ ⊢ x : Tensor[T, σ]
T ∈ {F16, BF16, F32, F64}              (only floating point; same set as softmax)
0 ≤ axis < |σ|                         (compile-time integer)
─────────────────────────────
Γ ⊢ log_softmax x axis : Tensor[T, σ]

Output shape, dtype, and placement equal the input. Runtime axes are not supported.

The bootstrap compiler’s softmax floating set is the frem dtypes (F16, BF16, F32, F8E4M3, F8E5M2). This section uses that same implemented set, not a wider IEEE set. F64 remains a spec-only dtype until the compiler grows it.

3.2 Operational meaning

For each fiber along axis:

m = max(x)
log_softmax(x) = (x - m) - log(sum(exp(x - m)))

The result is a log-probability vector (sums to 0 in the log domain after exp). Empty fibers are excluded by the existing “static dims ≥ 1” axiom.

3.3 Lowering

The user-facing form is one operation. The compiler decomposes it to Phase 0 primitives on both host and device (device softmax keeps its runtime kernel; log_softmax does not add a new FFI):

max_along     = reduce_max x axis  (keepdims)
shifted       = x - broadcast(max_along, shape(x))
exp_shifted   = exp(shifted)
sum_along     = reduce_sum exp_shifted axis  (keepdims)
log_sum       = log(broadcast(sum_along, shape(x)))
result        = shifted - log_sum

3.4 Surface

3.5 Tests

• Last-axis log_softmax on rank 2 and rank 3.
• Interior axis.
• Integer dtype rejected.
• Axis out of range rejected.
• Emitted MLIR is accepted by mlir-opt (golden).

4. take_along_axis

4.1 Typing rule

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

ρ ≡ σ \ axis is structural shape equality in this slice (same as other Phase 1 structural checks). Symbolic names must match exactly; there is no new SMT obligation.

The StoryLM-0 instance:

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

This is not 2D gather. gather stays embedding lookup.

4.2 Operational meaning

Write σ = α ++ [C] ++ β with |α| = axis. For every index tuple (a, b) into α and β,

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

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

4.3 Lowering (host)

1. A bounds linalg.generic over indices (compare against 0 and C).
2. A second linalg.generic over the output shape that tensor.extracts values at coordinates with axis replaced by the cast index.

Device placement is rejected at type-check until a checked device kernel exists (same stance as current gather).

4.4 Surface

Named parameter axis is required on the .kina form, matching softmax.

4.5 Tests

• Rank-3 last-axis select (the NLL pattern).
• Interior axis ([B, V, T] with axis=1).
• Shape mismatch rejected (indices not σ without axis).
• Non-I64 indices rejected.
• Device operands rejected.
• Axis out of range / rank-1 values rejected.
• Composed log_softmax + take_along_axis + neg example (storylm0_nll.kina).

5. Negative log-likelihood (derived)

nll = - values.log_softmax(axis).take_along_axis(targets, axis)

Mean or sum reduction is the existing reduce primitive. No nll_loss keyword in this slice.

6. AD

transform_ad rewrites these nodes in the forward walk (children only). VJP rules, the IR-internal scatter used by take_along_axis, and the missing neg VJP are in Section 10.

7. Open issues

1. Device take_along_axis. Follow-on; needs a checked kernel or a documented unchecked mode.
2. VJP rules. Specified in Section 10.
3. F64. Spec softmax lists F64; the bootstrap compiler does not. This section tracks the implemented floating set.
4. Index dim C symbolic. Bounds use tensor.dim at emit, same as gather’s vocab axis.