diff --git a/CHANGELOG.md b/CHANGELOG.md index 86d6f4b9..39e05341 100644 --- a/CHANGELOG.md +++ b/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` (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`; diff --git a/README.md b/README.md index 5af8b59d..bdd2438d 100644 --- a/README.md +++ b/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() --- -## 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 diff --git a/docs/antora.yml b/docs/antora.yml index 414e8a59..8691ea07 100644 --- a/docs/antora.yml +++ b/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 diff --git a/docs/modules/ROOT/partials/ops/tensorops/argmax.adoc b/docs/modules/ROOT/partials/ops/tensorops/argmax.adoc new file mode 100644 index 00000000..b98e4e25 --- /dev/null +++ b/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` — 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 = 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[] diff --git a/gradle.properties b/gradle.properties index 9e08b4f4..ab51de0e 100644 --- a/gradle.properties +++ b/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/