Phase 1, Section 17: Deterministic `argmax`

StoryLM-0 decode slice. Not a Phase 1 exit-criteria replacement. This section locks the Phase 0 argmax reduce op. It does not add a new token.

1. What this section commits to

Greedy decoding needs a deterministic class index from [B, T, V] logits: for each batch/position fiber, the index of the largest logit.

Phase 0 §8 already lists argmax (and argmin) among reduce ops and says the last two have output dtype I64 regardless of input. The typed-rule box still writes Tensor[dt, …] for every reduce. This section follows the prose, not the box: the result dtype is I64.

The surface already exists:

x.reduce(op=argmax, axis=2)
x.reduce_keep(op=argmax, axis=2)

No new keyword, IR constructor name, or FFI. axis is a compile-time integer, same as every other reduce. keepdims is the existing reduce-keep rule.

argmin stays unimplemented in this slice (same stance as prod).

2. Why lock this now

AGENTS.md step 4 asks for deterministic argmax after functional AdamW, then seeded sampling, then KV-cache. Sampling and cache semantics are still unspecified. This section only locks greedy index selection.

3. Typing

Admissibility is the same set as reduce max / reduce min:

Γ ⊢ x : Tensor[T, σ]
T ∈ frem ∪ {I64, I32, I16, I8, I4}
0 ≤ axis < |σ|
placement = Host                         (this slice)
─────────────────────────────────────────
Γ ⊢ reduce[argmax, axis] x : Tensor[I64, σ \ axis]

reduce-keep is the same rule with σ[axis] replaced by 1.

frem is the implemented floating set (F16, BF16, F32, F8E4M3, F8E5M2). Bool is rejected. Device placement is rejected at type-check until a checked device kernel exists (same stance as take_along_axis).

The StoryLM-0 instance:

logits : Tensor[F32, [1, 4, 8], @host]
axis   = 2
result : Tensor[I64, [1, 4], @host]

4. Operational meaning

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

out[a ++ b] = the smallest i in 0 .. C-1 such that fiber[i] = max(fiber)

Tie-break is lowest index. That matches NumPy and PyTorch argmax. The lowering uses a strict greater-than (arith.cmpf ogt on floats, arith.cmpi sgt on signed ints) so an equal value keeps the already-recorded index.

The numeric instance uses finite F32. NaN, ±∞-only fibers, and empty axes are unspecified here. Empty tensors are already excluded by the existing “static dims ≥ 1” axiom.

keepdims=true writes the same index into a size-1 axis.

5. Lowering (host)

A linalg.generic with two outputs over the input iterators (the reduced axis is "reduction"; the others are "parallel"):

1. Running max, same dtype as x, filled with the existing reduce max identity (−∞ for floats; signed-int minimum).
2. Running index, I64, filled with 0.

The body, for float x (float8 compares in f32 after extf, then truncf back for the max yield):

%gt      = arith.cmpf ogt, %x, %acc_max
%new_max = arith.select %gt, %x, %acc_max
%i       = linalg.index <axis> : index
%i64     = arith.index_cast %i : index to i64
%new_idx = arith.select %gt, %i64, %acc_idx
linalg.yield %new_max, %new_idx

The user-visible result is the index tensor only. Device placement fails emit if it reaches the GPU reduce FFI (no argmax opcode).

6. AD

argmax is discrete. There is no VJP in this slice. transform_ad fails with Argmax VJP is unsupported.

7. Tests

• argmax_axis1.kina — [2, 3] F32 → [2] I64; golden accepted by mlir-opt; emitted generic contains linalg.index.
• argmax_keep_axis1.kina — keepdims; result [2, 1] I64.
• storylm0_greedy.kina — §3 instance; wrapper fn result tensor<1x4xi64>.
• invalid_argmax_bool.kina — Bool rejected at type-check.
• invalid_argmax_device.kina — device placement rejected at type-check.
• scripts/validate_argmax.py compiles the greedy fn to a .so, injects patterned finite F32 logits including a tie, and matches NumPy argmax (lowest index).

8. Out of scope

• argmin, prod, any, all.
• Seeded sampling, temperature, top-k, top-p.
• KV-cache update semantics.
• Composing decoder forward + argmax as a generation loop (now Section 20).
• Device argmax.
• A VJP (including straight-through estimators).
• NaN policy.

9. Open issues

1. Seeded sampling. Specified as host Uniform + Gumbel-max in Section 18.
2. KV-cache / generate. Specified as one-step concat in Section 19 and a host greedy loop in Section 20.
3. Device argmax. Follow-on; needs a checked kernel or a documented unchecked mode.
4. argmin. Same lowering with strict less-than; not this slice.
5. NaN. Unspecified. Do not guess a payload.