fix(dag): infer argMax output spec as reduced i32 (fixes #876)

[michalharakal]

The dag{} builder path never inferred argMax's output spec. GraphDsl.inferDagOutputSpecs has cases for
reductions / reshape / matmul / concat / conv but not argMax, so it fell back to echoing operand-0's
shape + dtype. An argMax over f32 logits thus recorded an f32, full-shape output; the (correct)
StableHloConverter then faithfully emitted invalid IR:

• stablehlo.constant dense<N> : tensor<…xf32> — an integer literal for an f32 tensor, which
  iree-compile rejects ("unexpected decimal integer literal for a floating point value");
• a final stablehlo.reduce whose result kept the reduced dim (1x4x8, not 1x4).

The VoidTensorOps path already inferred this correctly (reduced + Int32), which is why the FunctionGemma
NN-DSL export was board-verified and only the raw dag{} path broke — surfaced by the functiongemma-270m
conformance row (skainet-iree-conformance #24).

Fix

Add an argmax/argmin case to inferDagOutputSpecs — reduced shape (like sum/mean) with an Int32
index dtype — and let the local spec() helper take a dtype override. No converter change — it was
correct given a correct spec.

Test

ArgMaxOperationsConverterTest.testArgMaxThroughDagBuilderProducesCompilableIr traces dag { argMax } end to
end and asserts the i32 sentinel + collapsed reduce + no int-literal-f32 constant — the exact conditions #876
broke. The pre-existing converter test used a hand-built correct node, bypassing the inference (why the bug
slipped). Verified: skainet-compile-hlo argMax tests 5/5, full skainet-lang-dag suite green.

Fixes #876.

[github-actions]

📖 Documentation Preview

The documentation has been built successfully for this PR.

Generated Files:

• Operator documentation: docs/modules/operators/_generated_/
• JSON schema output: operators.json

Artifacts:

• Download the documentation-preview-878 artifact to view the complete documentation locally.

This comment will be updated automatically when the PR is updated.