release: 0.35.0 — argMax op (version, changelog, README, op docs)

CHANGELOG.md

@@ -2,7 +2,21 @@
 
 ## [Unreleased]
 
-### Changed
+## [0.35.0] - 2026-07-08
+
+### Added
+
+- **`argMax(tensor, dim)` tensor op.** Index of the maximum value along a dimension, with ties
+  resolved to the **lowest** index (numpy/greedy); the reduced dimension is removed (no keepdim).
+  Non-differentiable by design (index selection). Like `scaledDotProductAttention`, it stays a
+  single op and is lowered at the **StableHLO stage** rather than needing a new primitive:
+  `ArgMaxOperationsConverter` composes `iota` + reduce-`maximum` + `broadcast_in_dim` + `compare EQ`
+  + `select` + reduce-`minimum` (the same `stablehlo` primitives the causal-mask code already
+  emits) — so no variadic reducer region and no new DSL ops. The eager CPU kernel is a scalar
+  reduction-to-index. Indices are `i32` in the compiled StableHLO; the **eager** path materializes
+  them as index-valued floats in the input dtype so the result is a portable `Tensor<T, V>` (an i32
+  payload inside a float tensor is unreadable on Kotlin/Native + Wasm). Unblocks folding an LLM's
+  `logits → token-ids` argmax tail into the DSL trace instead of a post-hoc MLIR rewrite. (PR #800)
 
 - **BREAKING (coordinates): `skainet-data-simple` now publishes under its module name.** The artifactId
   changes from the mismatched `sk.ainet.core:skainet-data-basic` to `sk.ainet.core:skainet-data-simple`;

README.md

@@ -40,7 +40,7 @@ Add the core dependencies (Gradle Kotlin DSL):
 ```kotlin
 dependencies {
     // Recommended: import the umbrella BOM and drop versions on the engine modules.
-    implementation(platform("sk.ainet:skainet-bom:0.34.0"))
+    implementation(platform("sk.ainet:skainet-bom:0.35.0"))
 
     implementation("sk.ainet.core:skainet-lang-core")
     implementation("sk.ainet.core:skainet-backend-cpu")
@@ -284,8 +284,9 @@ val withoutLabel = dataPipeline<RawDataset>()
 
 ---
 
-## What's New in 0.34.0
+## What's New in 0.35.0
 
+- **`argMax(dim)` tensor op** — index of the maximum along a dimension (ties → lowest index). Lowered to StableHLO by composing `iota` + reduce-max + `compare`/`select` + reduce-min — a single op, no new primitive — plus an eager CPU kernel. Lets an LLM's `logits → token-ids` argmax tail live inside the DSL trace.
 - **URI-backed data sources** — new `skainet-data-source` module: `file://`, `https://`, and Hugging Face URIs, raw-format parsers (CSV/TSV/JSON/JSONL), suspendable data pipelines
 - **Dataset views and richer batches** — seeded shuffle, stratified split, filter views, batch/epoch flows, batch indices + metadata
 - **bf16-native DSL → StableHLO export** — weights reach the matmuls as bf16, verified down to an aarch64 vmfb

docs/antora.yml

@@ -15,7 +15,7 @@ asciidoc:
     framework_name: SKaiNET
     # Current SKaiNET release — bump once per release; referenced as
     # {skainet_version} in dependency snippets (blocks need subs="attributes+").
-    skainet_version: 0.34.0
+    skainet_version: 0.35.0
     ksp_version: 2.2.21-2.0.5
     dokka_version: 2.1.0
     asciidoctorj_version: 3.0.0

docs/modules/ROOT/partials/ops/tensorops/argmax.adoc

@@ -0,0 +1,47 @@
+// tag::math[]
+Given a tensor stem:[X] and a reduction dimension stem:[d] of size stem:[n], `argMax` returns, for
+each position of the remaining dimensions, the index stem:[k^\*] of the maximum value along stem:[d]:
+
+[stem]
+++++
+k^\* = \min \{\, k \in [0, n) : X_{\ldots, k, \ldots} = \max_{0 \le j < n} X_{\ldots, j, \ldots} \,\}
+++++
+
+Ties resolve to the *lowest* index (the `\min` above — numpy/greedy semantics). Dimension stem:[d]
+is removed from the output (no keepdim), and the result holds integer indices.
+// end::math[]
+
+// tag::intuition[]
+`argMax` turns scores into a decision: "which entry won?" It is the final step of greedy decoding
+(pick the highest-probability next token from a `[..., vocab]` logits tensor) and of classification
+(pick the top class). Because it selects an index rather than blending values, it is
+*non-differentiable* and carries no backward rule.
+
+SKaiNET keeps `argMax` a *single* op (like `scaledDotProductAttention`) and lowers it at the
+StableHLO stage rather than adding a dedicated primitive. `ArgMaxOperationsConverter` composes it
+from ops the exporter already emits:
+
+. `stablehlo.iota` along stem:[d] — the candidate indices,
+. reduce-`maximum` along stem:[d], then `broadcast_in_dim` back — the per-position max value,
+. `compare EQ` of the input against that max — a boolean mask of the maxima,
+. `select` mask ? index : stem:[n] (an out-of-range sentinel), then reduce-`minimum` — the lowest
+  index among the maxima.
+
+*Output dtype.* Indices are `i32` in the compiled StableHLO. The *eager* CPU path materializes them
+as index-valued floats in the input dtype, so the eager result is a portable `Tensor<T, V>` — an
+`i32` payload inside a float tensor is unreadable on Kotlin/Native and Wasm.
+// end::intuition[]
+
+// tag::examples[]
+[source,kotlin]
+----
+// [batch, vocab] logits -> [batch] token ids (greedy decode)
+val logits: Tensor<FP32, FloatArray> = tensor(shape(1, 4)) { float(0.1f, 3.2f, 0.5f, 2.0f) }
+val ids = ops.argMax(logits, dim = -1)   // shape(1); value 1 (the max, 3.2, is at index 1)
+----
+// end::examples[]
+
+// tag::references[]
+* https://numpy.org/doc/stable/reference/generated/numpy.argmax.html[numpy.argmax] — the tie-breaking (lowest index) reference semantics.
+* https://openxla.org/stablehlo/spec#reduce[StableHLO `reduce`] and https://openxla.org/stablehlo/spec#iota[`iota`] — the primitives the lowering composes.
+// end::references[]

gradle.properties

@@ -1,5 +1,5 @@
 GROUP=sk.ainet.core
-VERSION_NAME=0.34.0
+VERSION_NAME=0.35.0
 POM_DESCRIPTION=SKaiNET
 
 POM_URL=https://github.com/SKaiNET-developers/skainet/