From 74cfabc8cccff747348f6f60fdfc4631eb8f137c Mon Sep 17 00:00:00 2001 From: michal harakal Date: Wed, 8 Jul 2026 15:20:58 +0200 Subject: [PATCH 1/3] feat(ops): add argMax(dim) tensor op with StableHLO lowering + CPU kernel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit argMax is a real, HW-agnostic DSL op — the index of the max along a dim (ties -> lowest index), returning Int32 indices. Non-differentiable. Like scaledDotProductAttention it stays a single op and is lowered at the StableHLO stage rather than having a dedicated hardware primitive: ArgMaxOperationsConverter composes iota + reduce-max + broadcast + compare + select + reduce-min (the same stablehlo primitives the causal-mask code already emits), so no variadic reducer region is needed. Registered in StableHloConverterFactory (basic/extended/fast). Eager path is a scalar reduction-to-index kernel in DefaultCpuOps. Shape floor in VoidTensorOps; delegation added to RecordingTensorOpsDecorator and the dag TestTensorOps. Enables the FunctionGemma compiled export to fold its logits->token-ids tail into the DSL trace (retiring an external Python MLIR rewrite). Tested: eager numeric (last-dim, ties, middle-dim) + StableHLO lowering (the [1,24,262153]->[1,24] gemma tail); full compile-hlo suite green. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 37 ++++++ .../tensor/ops/DefaultCpuOpsArgMaxTest.kt | 58 +++++++++ .../sk/ainet/tape/RecordingExecution.kt | 1 + .../compile/graph/ComputeGraphExecutorTest.kt | 1 + .../compile/hlo/StableHloConverterFactory.kt | 7 ++ .../converters/ArgMaxOperationsConverter.kt | 114 ++++++++++++++++++ .../hlo/ArgMaxOperationsConverterTest.kt | 85 +++++++++++++ .../sk/ainet/lang/tensor/ops/TensorOps.kt | 12 ++ .../sk/ainet/lang/tensor/ops/VoidTensorOps.kt | 10 ++ 9 files changed, 325 insertions(+) create mode 100644 skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsArgMaxTest.kt create mode 100644 skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ArgMaxOperationsConverter.kt create mode 100644 skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ArgMaxOperationsConverterTest.kt diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt index 8aca6dc2..7070b21e 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt @@ -2050,6 +2050,43 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory } } + @TensorOp() + @InProgress("cpu", owner = "team:cpu", issue = "task-ops.md#op-argmax") + override fun argMax(tensor: Tensor, dim: Int): Tensor { + val rank = tensor.rank + require(rank >= 1) { "argMax: input must have rank >= 1" } + val nd = if (dim < 0) dim + rank else dim + require(nd in 0 until rank) { "argMax: dim $dim out of range for rank $rank" } + val inDims = tensor.shape.dimensions + // Reduced (dim removed) shape; scalar -> Shape(1) to match VoidTensorOps. + val reduced = IntArray(rank - 1) + run { var o = 0; for (i in 0 until rank) if (i != nd) reduced[o++] = inDims[i] } + val outShape = if (reduced.isEmpty()) Shape(1) else Shape(reduced) + val outCount = reduced.fold(1) { a, b -> a * b } + val indices = IntArray(outCount) + val outIdx = IntArray(reduced.size) + for (flat in 0 until outCount) { + val inIdx = IntArray(rank) + run { var o = 0; for (i in 0 until rank) if (i != nd) inIdx[i] = outIdx[o++] } + var bestK = 0 + var bestV = Float.NEGATIVE_INFINITY + for (k in 0 until inDims[nd]) { + inIdx[nd] = k + val v = (tensor.data.get(*inIdx) as Number).toFloat() + if (v > bestV) { bestV = v; bestK = k } // strict > keeps the LOWEST index on ties + } + indices[flat] = bestK + var d = reduced.size - 1 + while (d >= 0) { outIdx[d]++; if (outIdx[d] < reduced[d]) break; outIdx[d] = 0; d-- } + } + @Suppress("UNCHECKED_CAST") + val int32 = sk.ainet.lang.types.Int32::class as KClass + @Suppress("UNCHECKED_CAST") + val outData = dataFactory.fromIntArray(outShape, int32, indices) + as sk.ainet.lang.tensor.data.TensorData + return CpuTensor(outData, this, int32, GradState(requiresGrad = false)) + } + @TensorOp() @InProgress("cpu", owner = "team:cpu", issue = "task-ops.md#op-mean") override fun mean( diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsArgMaxTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsArgMaxTest.kt new file mode 100644 index 00000000..fc8b48e3 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsArgMaxTest.kt @@ -0,0 +1,58 @@ +package sk.ainet.sk.ainet.exec.tensor.ops + +import kotlin.test.Test +import kotlin.test.assertEquals +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.data +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.dsl.tensor +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Int32 + +class DefaultCpuOpsArgMaxTest { + private val ctx = DirectCpuExecutionContext() + private val ops get() = ctx.ops + + // argMax returns an Int32-backed tensor (typed by the op signature); + // read the raw index via the underlying TensorData. + private fun idx(t: Tensor, vararg i: Int): Int = (t.data.get(*i) as Number).toInt() + + @Test + fun argmax_last_dim_returns_indices_drops_dim_and_is_int32() { + data(ctx) { _ -> + // [2,3]: row0 max at col2 (7), row1 max at col0 (9) + val v = floatArrayOf(1f, 3f, 7f, 9f, 2f, 5f) + val t = tensor { shape(2, 3) { init { v[it[0] * 3 + it[1]] } } } + val r = ops.argMax(t, dim = -1) + assertEquals(Shape(2), r.shape) + assertEquals(Int32::class, r.dtype) // runtime dtype is Int32 despite the static type + assertEquals(2, idx(r, 0)) + assertEquals(0, idx(r, 1)) + } + } + + @Test + fun argmax_ties_break_to_lowest_index() { + data(ctx) { _ -> + val v = floatArrayOf(5f, 5f, 2f, 5f) // maxima at indices 0,1,3 -> lowest is 0 + val t = tensor { shape(1, 4) { init { v[it[1]] } } } + val r = ops.argMax(t, dim = 1) + assertEquals(Shape(1), r.shape) + assertEquals(0, idx(r, 0)) + } + } + + @Test + fun argmax_middle_dim() { + data(ctx) { _ -> + // [2,3,1]: reduce dim 1 -> [2,1]; per (row, last) the max over the 3 middle entries + val v = floatArrayOf(/*r0*/ 4f, 9f, 1f, /*r1*/ 2f, 0f, 8f) + val t = tensor { shape(2, 3, 1) { init { v[it[0] * 3 + it[1]] } } } + val r = ops.argMax(t, dim = 1) + assertEquals(Shape(2, 1), r.shape) + assertEquals(1, idx(r, 0, 0)) // row0 max 9 at mid=1 + assertEquals(2, idx(r, 1, 0)) // row1 max 8 at mid=2 + } + } +} diff --git a/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt b/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt index 2a32e39e..78eb5321 100644 --- a/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt +++ b/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt @@ -447,6 +447,7 @@ internal class RecordingTensorOpsDecorator(private val base: TensorOps) : Tensor override fun sum(tensor: Tensor, dim: Int?): Tensor = base.sum(tensor, dim) override fun mean(tensor: Tensor, dim: Int?): Tensor = base.mean(tensor, dim) override fun variance(tensor: Tensor, dim: Int?): Tensor = base.variance(tensor, dim) + override fun argMax(tensor: Tensor, dim: Int): Tensor = base.argMax(tensor, dim) override fun sqrt(tensor: Tensor): Tensor = base.sqrt(tensor) override fun log(tensor: Tensor): Tensor = base.log(tensor) override fun log2(tensor: Tensor): Tensor = base.log2(tensor) diff --git a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt index b5ea2d13..6ac6a282 100644 --- a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt +++ b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt @@ -177,6 +177,7 @@ private class TestTensorOps : TensorOps { override fun sum(tensor: Tensor, dim: Int?): Tensor = tensor override fun mean(tensor: Tensor, dim: Int?): Tensor = tensor override fun variance(tensor: Tensor, dim: Int?): Tensor = tensor + override fun argMax(tensor: Tensor, dim: Int): Tensor = tensor override fun sqrt(tensor: Tensor): Tensor = tensor override fun pow(a: Tensor, b: Tensor): Tensor = a override fun powScalar(a: Tensor, n: Number): Tensor = a diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverterFactory.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverterFactory.kt index 22e04676..726c0ebe 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverterFactory.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverterFactory.kt @@ -1,6 +1,7 @@ package sk.ainet.compile.hlo import sk.ainet.compile.hlo.converters.ActivationOperationsConverter +import sk.ainet.compile.hlo.converters.ArgMaxOperationsConverter import sk.ainet.compile.hlo.converters.AttentionOperationsConverter import sk.ainet.compile.hlo.converters.ConstantOperationsConverter import sk.ainet.compile.hlo.converters.GatherOperationsConverter @@ -57,6 +58,8 @@ public object StableHloConverterFactory { // Register reduction operations converter registry.register(ReductionOperationsConverter()) + // argMax: logits -> index, lowered to reduce-max + broadcast + compare + iota + select + reduce-min + registry.register(ArgMaxOperationsConverter()) // Register elementwise unary math converter (sqrt, exp, log, abs, …). // Must be present so downstream consumers don't cascade-fail with @@ -114,6 +117,8 @@ public object StableHloConverterFactory { // Register reduction operations converter registry.register(ReductionOperationsConverter()) + // argMax: logits -> index, lowered to reduce-max + broadcast + compare + iota + select + reduce-min + registry.register(ArgMaxOperationsConverter()) // Register elementwise unary math converter (sqrt, exp, log, abs, …). // Must be present so downstream consumers don't cascade-fail with @@ -157,6 +162,8 @@ public object StableHloConverterFactory { registry.register(ActivationOperationsConverter()) registry.register(ShapeOperationsConverter()) registry.register(ReductionOperationsConverter()) + // argMax: logits -> index, lowered to reduce-max + broadcast + compare + iota + select + reduce-min + registry.register(ArgMaxOperationsConverter()) registry.register(UnaryMathConverter()) registry.register(ScalarOperationsConverter()) registry.register(ConstantOperationsConverter()) diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ArgMaxOperationsConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ArgMaxOperationsConverter.kt new file mode 100644 index 00000000..34a8dfcd --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ArgMaxOperationsConverter.kt @@ -0,0 +1,114 @@ +package sk.ainet.compile.hlo.converters + +import sk.ainet.compile.hlo.ConversionContext +import sk.ainet.compile.hlo.ConversionResult +import sk.ainet.compile.hlo.StableHloOperationConverter +import sk.ainet.lang.graph.GraphNode + +/** + * Converter for `argMax` — the index of the maximum value along a dimension, + * with ties resolved to the LOWEST index (numpy/greedy semantics). The reduced + * dimension is removed from the output; the result is an `i32` index tensor. + * + * StableHLO has no argmax primitive, so — like softmax and the attention mask — + * it is lowered by COMPOSING single-op primitives that this codebase already + * emits (no variadic `stablehlo.reduce` with a custom reducer region, which the + * MLIR validator here has no precedent for): + * + * ``` + * maxV = reduce(x) applies stablehlo.maximum across [dim] // per-position max value + * maxB = broadcast_in_dim maxV -> input shape + * isMax = compare EQ, x, maxB // i1 mask of maxima + * idx = iota dim = dim // i32 indices along dim + * cand = select isMax, idx, // non-maxima -> out-of-range + * argIdx = reduce(cand) applies stablehlo.minimum across [dim] // lowest index of the max + * ``` + * + * The `compare EQ` is exact: `maxB` holds bit-identical copies of actual input + * values, so the max element(s) compare equal; `minimum` then picks the lowest + * index among ties and ignores non-maxima (masked to the `dimSize` sentinel, + * which exceeds every valid `0..dimSize-1` index). + */ +public class ArgMaxOperationsConverter : StableHloOperationConverter { + + override val supportedOperations: Set = setOf("argMax", "argmax") + + override fun convert( + node: GraphNode, + operands: List, + context: ConversionContext + ): ConversionResult { + if (operands.size != 1) { + return ConversionResult.Failure( + "argMax requires exactly 1 operand, got ${operands.size}", + "Unsupported argMax arity for node ${node.id}" + ) + } + val inputSpec = node.inputs.firstOrNull() + ?: return ConversionResult.Failure("argMax: missing input spec", "node ${node.id}") + val outputSpec = node.outputs.firstOrNull() + ?: return ConversionResult.Failure("argMax: missing output spec", "node ${node.id}") + val inShape = inputSpec.shape + ?: return ConversionResult.Failure("argMax: static input shape required", "node ${node.id}") + + val rank = inShape.size + val dimParam = (node.operation.parameters["dim"] as? Int) + ?: (node.operation.parameters["axis"] as? Int) + ?: (rank - 1) + val nd = if (dimParam < 0) rank + dimParam else dimParam + if (nd !in 0 until rank) { + return ConversionResult.Failure("argMax: dim $dimParam out of range for rank $rank", "node ${node.id}") + } + + val tm = context.getTypeMapper() + val valElem = tm.mapDType(inputSpec.dtype) // e.g. "f32" + val idxElem = tm.mapDType(outputSpec.dtype) // "i32" + fun typeOf(dims: List, elem: String): String = + if (dims.isEmpty()) "tensor<$elem>" else "tensor<${dims.joinToString("x")}x$elem>" + + val fullValType = typeOf(inShape, valElem) // tensor<1x24x262153xf32> + val fullIdxType = typeOf(inShape, idxElem) // tensor<1x24x262153xi32> + val fullI1Type = typeOf(inShape, "i1") // tensor<1x24x262153xi1> + val reducedDims = inShape.filterIndexed { i, _ -> i != nd } + val reducedValType = typeOf(reducedDims, valElem) // tensor<1x24xf32> + val idxOutType = tm.mapTensorType(outputSpec) // tensor<1x24xi32> + val broadcastDims = (0 until rank).filter { it != nd }.joinToString(", ") + val dimSize = inShape[nd] // sentinel: > any valid index + val negInf = tm.negInfBits(valElem) + + val ops = mutableListOf() + fun emit(op: String) { ops += op; context.emitOperation(op) } + + val maxInit = context.nextTempValue() + emit("$maxInit = stablehlo.constant dense<$negInf> : tensor<$valElem>") + val maxV = context.nextTempValue() + emit( + "$maxV = stablehlo.reduce(${operands[0]} init: $maxInit) " + + "applies stablehlo.maximum across dimensions = [$nd] : " + + "($fullValType, tensor<$valElem>) -> $reducedValType" + ) + val maxB = context.nextTempValue() + emit( + "$maxB = stablehlo.broadcast_in_dim $maxV, dims = [$broadcastDims] : " + + "($reducedValType) -> $fullValType" + ) + val isMax = context.nextTempValue() + emit("$isMax = stablehlo.compare EQ, ${operands[0]}, $maxB : ($fullValType, $fullValType) -> $fullI1Type") + val idx = context.nextTempValue() + emit("$idx = stablehlo.iota dim = $nd : $fullIdxType") + val sentinel = context.nextTempValue() + emit("$sentinel = stablehlo.constant dense<$dimSize> : $fullIdxType") + val cand = context.nextTempValue() + emit("$cand = stablehlo.select $isMax, $idx, $sentinel : $fullI1Type, $fullIdxType") + val minInit = context.nextTempValue() + emit("$minInit = stablehlo.constant dense<$dimSize> : tensor<$idxElem>") + val argIdx = context.nextTempValue() + emit( + "$argIdx = stablehlo.reduce($cand init: $minInit) " + + "applies stablehlo.minimum across dimensions = [$nd] : " + + "($fullIdxType, tensor<$idxElem>) -> $idxOutType" + ) + + return ConversionResult.Success(outputValueName = argIdx, emittedOperations = ops) + } +} diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ArgMaxOperationsConverterTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ArgMaxOperationsConverterTest.kt new file mode 100644 index 00000000..31f1b7ed --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ArgMaxOperationsConverterTest.kt @@ -0,0 +1,85 @@ +package sk.ainet.compile.hlo + +import sk.ainet.compile.hlo.converters.ArgMaxOperationsConverter +import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.ops.Operation +import sk.ainet.lang.tensor.ops.TensorSpec +import sk.ainet.lang.tensor.ops.ValidationResult +import sk.ainet.lang.types.DType +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertIs +import kotlin.test.assertTrue + +class ArgMaxOperationsConverterTest { + + private val converter = ArgMaxOperationsConverter() + private val typeMapper = TypeMapper() + private val context = ConversionContext(typeMapper) + + @Test + fun testSupportedOperations() { + assertEquals(setOf("argMax", "argmax"), converter.supportedOperations) + } + + @Test + fun testRegistryIntegration() { + // The default factory registers it (createBasic); assert via a fresh registry mirroring it. + val registry = StableHloOperationRegistry() + registry.register(ArgMaxOperationsConverter()) + assertTrue(registry.isSupported("argMax")) + assertTrue(registry.isSupported("argmax")) + } + + @Test + fun testArgMaxLoweringGemmaTail() { + // The FunctionGemma tail: [1, 24, 262153] f32 logits -> [1, 24] i32 token ids, dim = -1. + val operation = createMockOperation("argMax", mapOf("dim" to -1)) + val inputSpec = TensorSpec("logits", listOf(1, 24, 262153), "FP32") + val outputSpec = TensorSpec("ids", listOf(1, 24), "Int32") + val node = GraphNode("test_argmax", operation, listOf(inputSpec), listOf(outputSpec)) + + val result = converter.convert(node, listOf("%logits"), context) + + assertIs(result) + assertTrue(result.outputValueName.startsWith("%v")) + val ops = result.emittedOperations + // Composed from single-op primitives (no variadic reducer region). + assertTrue(ops.any { it.contains("stablehlo.iota dim = 2") }, "iota over the reduced dim") + assertTrue(ops.any { it.contains("stablehlo.reduce(") && it.contains("stablehlo.maximum") }, "reduce-max") + assertTrue(ops.any { it.contains("stablehlo.broadcast_in_dim") }, "broadcast max back") + assertTrue(ops.any { it.contains("stablehlo.compare EQ") }, "equality mask of maxima") + assertTrue(ops.any { it.contains("stablehlo.select") }, "select index or sentinel") + assertTrue(ops.any { it.contains("stablehlo.reduce(") && it.contains("stablehlo.minimum") }, "reduce-min -> lowest index") + // Sentinel = dim size (262153) so non-maxima lose the min; output is i32. + assertTrue(ops.any { it.contains("dense<262153>") }, "out-of-range sentinel is the dim size") + assertTrue(ops.any { it.contains("-> tensor<1x24xi32>") }, "output is the i32 index tensor") + } + + @Test + fun testInvalidOperandCount() { + val operation = createMockOperation("argMax", mapOf("dim" to 1)) + val outputSpec = TensorSpec("ids", listOf(2), "Int32") + val inputSpec = TensorSpec("x", listOf(2, 3), "FP32") + val node = GraphNode("test_argmax", operation, listOf(inputSpec), listOf(outputSpec)) + + val resultZero = converter.convert(node, emptyList(), context) + assertIs(resultZero) + assertTrue(resultZero.error.contains("requires exactly 1 operand")) + } + + private fun createMockOperation(name: String, parameters: Map): Operation { + return object : Operation { + override val name: String = name + override val type: String = "reduction" + override val parameters: Map = parameters + override fun execute(inputs: List>): List> = + throw UnsupportedOperationException("Mock operation") + override fun validateInputs(inputs: List): ValidationResult = ValidationResult.Valid + override fun inferOutputs(inputs: List): List = emptyList() + override fun clone(newParameters: Map): Operation = createMockOperation(name, newParameters) + override fun serialize(): Map = mapOf("name" to name, "type" to type, "parameters" to parameters) + } + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt index d326cd2a..7f54239d 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt @@ -193,6 +193,18 @@ public interface TensorOps { @Diff public fun variance(tensor: Tensor, dim: Int? = null): Tensor + /** + * Index of the maximum value along [dim]. The reduced dimension is removed from + * the output (no keepdim); ties resolve to the LOWEST index (numpy/greedy + * semantics). The output holds `Int32` indices, so downstream consumers should + * treat it as an index tensor. Non-differentiable by design (index selection). + * + * Like [scaledDotProductAttention], this stays a single op and is lowered at the + * StableHLO stage (`ArgMaxOperationsConverter`: reduce-max → broadcast → compare → + * iota → select → reduce-min) rather than having a dedicated hardware kernel. + */ + public fun argMax(tensor: Tensor, dim: Int): Tensor + // Mathematical functions @Diff public fun sqrt(tensor: Tensor): Tensor diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt index 92dad4da..44f4467e 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt @@ -366,6 +366,16 @@ public class VoidTensorOps : TensorOps { return VoidOpsTensor(resultData, tensor.dtype) } + override fun argMax(tensor: Tensor, dim: Int): Tensor { + // Reduced shape (dim removed) with an Int32 index dtype — mirrors `convert`'s + // dtype-changing pattern (the traced graph node then carries i32 output). + val resultShape = calculateReductionShape(tensor.shape, dim, "argMax") + @Suppress("UNCHECKED_CAST") + val int32 = sk.ainet.lang.types.Int32::class as kotlin.reflect.KClass + val resultData = dataFactory.zeros(resultShape, int32) + return VoidOpsTensor(resultData, int32) + } + override fun sqrt(tensor: Tensor): Tensor { // Square root function preserves shape val resultData = dataFactory.zeros(tensor.shape, tensor.dtype) From 8938b8a020992cf3b0fe22f50a4db5c0460be73a Mon Sep 17 00:00:00 2001 From: michal harakal Date: Wed, 8 Jul 2026 16:25:48 +0200 Subject: [PATCH 2/3] chore(api): sync .api dumps for argMax (fixes apiCheck) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `./gradlew apiDump` after adding the public `argMax` op. Updates the jvm ABI dumps of every module that re-exposes it via `TensorOps` (interface method + DSL/graph re-exports): lang-core, lang-dag, backend-cpu, compile-hlo, compile-dag, compile-opt. Also captures pre-existing untracked public API in compile-opt/compile-hlo (FusedOpAllowList / OpGranularityPolicy / a converter ctor) that had drifted out of the golden dumps on develop — apiDump is generated from the actual code, so this only reconciles them. Verified: apiCheck green; js/wasmJs/wasmWasi compile clean; yarn.lock unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../api/jvm/skainet-backend-cpu.api | 1 + .../api/jvm/skainet-compile-dag.api | 11 +++++ .../api/jvm/skainet-compile-hlo.api | 29 +++++++++---- .../api/jvm/skainet-compile-opt.api | 42 +++++++++++++++++++ .../api/jvm/skainet-lang-core.api | 3 ++ .../api/jvm/skainet-lang-dag.api | 2 + 6 files changed, 81 insertions(+), 7 deletions(-) diff --git a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api index 9c6883d8..0ed4e5fb 100644 --- a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api +++ b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api @@ -208,6 +208,7 @@ public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/lang/tensor/o public fun abs (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun add (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun addScalar (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Number;)Lsk/ainet/lang/tensor/Tensor; + public fun argMax (Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor; public fun avgPool2d (Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;Z)Lsk/ainet/lang/tensor/Tensor; protected final fun broadcastShapes (Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/tensor/Shape;)Lsk/ainet/lang/tensor/Shape; protected final fun chooseQuantizedMatmulHeap (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; diff --git a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api index 9bd0f858..f98145a5 100644 --- a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api +++ b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api @@ -1,3 +1,14 @@ +public final class sk/ainet/compile/target/FusedOpAllowList : sk/ainet/compile/target/OpGranularityPolicy { + public fun (Ljava/lang/String;Ljava/util/Set;)V + public fun getTarget ()Ljava/lang/String; + public fun keepFused (Ljava/lang/String;)Z +} + +public abstract interface class sk/ainet/compile/target/OpGranularityPolicy { + public abstract fun getTarget ()Ljava/lang/String; + public abstract fun keepFused (Ljava/lang/String;)Z +} + public abstract interface class sk/ainet/lang/graph/ComputeGraph { public abstract fun addEdge (Lsk/ainet/lang/graph/GraphEdge;)Lsk/ainet/lang/graph/GraphEdge; public abstract fun addNode (Lsk/ainet/lang/graph/GraphNode;)Lsk/ainet/lang/graph/GraphNode; diff --git a/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api b/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api index ffbaa568..4cfb64cd 100644 --- a/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api +++ b/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api @@ -45,7 +45,9 @@ public final class sk/ainet/compile/hlo/ConversionContext { public fun (Lsk/ainet/compile/hlo/TypeMapper;)V public fun (Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/lang/graph/ComputeGraph;)V public fun (Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;)V - public synthetic fun (Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;Ljava/lang/String;)V + public fun (Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;Ljava/lang/String;Lsk/ainet/compile/target/OpGranularityPolicy;)V + public synthetic fun (Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;Ljava/lang/String;Lsk/ainet/compile/target/OpGranularityPolicy;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun clear ()V public final fun emitComment (Ljava/lang/String;)V public final fun emitEncodingAnnotation (Ljava/lang/String;ILsk/ainet/lang/tensor/ops/TensorSpec;)V @@ -54,9 +56,11 @@ public final class sk/ainet/compile/hlo/ConversionContext { public final fun emitOperation (Ljava/lang/String;)V public final fun getContent ()Ljava/lang/String; public final fun getExternalParameters ()Ljava/util/List; + public final fun getGranularity ()Lsk/ainet/compile/target/OpGranularityPolicy; public final fun getInputNodes (Lsk/ainet/lang/graph/GraphNode;)Ljava/util/List; public final fun getMaterializationPolicy ()Lsk/ainet/compile/hlo/ConstantMaterializationPolicy; public final fun getModuleDeclarations ()Ljava/lang/String; + public final fun getTarget ()Ljava/lang/String; public final fun getTypeMapper ()Lsk/ainet/compile/hlo/TypeMapper; public final fun getValueName (Ljava/lang/String;)Ljava/lang/String; public final fun getValueName (Ljava/lang/String;I)Ljava/lang/String; @@ -115,10 +119,10 @@ public final class sk/ainet/compile/hlo/ConversionResult$Unsupported : sk/ainet/ } public final class sk/ainet/compile/hlo/Dag2hloKt { - public static final fun toStableHlo (Lsk/ainet/lang/graph/ComputeGraph;Ljava/lang/String;)Lsk/ainet/compile/hlo/StableHloModule; - public static final fun toStableHlo (Lsk/ainet/lang/graph/ResolvedComputeGraph;Ljava/lang/String;Z)Lsk/ainet/compile/hlo/StableHloModule; - public static synthetic fun toStableHlo$default (Lsk/ainet/lang/graph/ComputeGraph;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/compile/hlo/StableHloModule; - public static synthetic fun toStableHlo$default (Lsk/ainet/lang/graph/ResolvedComputeGraph;Ljava/lang/String;ZILjava/lang/Object;)Lsk/ainet/compile/hlo/StableHloModule; + public static final fun toStableHlo (Lsk/ainet/lang/graph/ComputeGraph;Ljava/lang/String;Ljava/lang/String;Lsk/ainet/compile/target/OpGranularityPolicy;)Lsk/ainet/compile/hlo/StableHloModule; + public static final fun toStableHlo (Lsk/ainet/lang/graph/ResolvedComputeGraph;Ljava/lang/String;ZLjava/lang/String;Lsk/ainet/compile/target/OpGranularityPolicy;)Lsk/ainet/compile/hlo/StableHloModule; + public static synthetic fun toStableHlo$default (Lsk/ainet/lang/graph/ComputeGraph;Ljava/lang/String;Ljava/lang/String;Lsk/ainet/compile/target/OpGranularityPolicy;ILjava/lang/Object;)Lsk/ainet/compile/hlo/StableHloModule; + public static synthetic fun toStableHlo$default (Lsk/ainet/lang/graph/ResolvedComputeGraph;Ljava/lang/String;ZLjava/lang/String;Lsk/ainet/compile/target/OpGranularityPolicy;ILjava/lang/Object;)Lsk/ainet/compile/hlo/StableHloModule; public static final fun toStableHloLegacy (Lsk/ainet/lang/graph/ComputeGraph;Ljava/lang/String;)Lsk/ainet/compile/hlo/StableHloModule; public static synthetic fun toStableHloLegacy$default (Lsk/ainet/lang/graph/ComputeGraph;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/compile/hlo/StableHloModule; } @@ -180,7 +184,9 @@ public final class sk/ainet/compile/hlo/StableHloConverter { public fun (Lsk/ainet/compile/hlo/StableHloOperationRegistry;Lsk/ainet/compile/hlo/TypeMapper;)V public fun (Lsk/ainet/compile/hlo/StableHloOperationRegistry;Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/compile/hlo/MlirValidator;)V public fun (Lsk/ainet/compile/hlo/StableHloOperationRegistry;Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/compile/hlo/MlirValidator;Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;)V - public synthetic fun (Lsk/ainet/compile/hlo/StableHloOperationRegistry;Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/compile/hlo/MlirValidator;Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Lsk/ainet/compile/hlo/StableHloOperationRegistry;Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/compile/hlo/MlirValidator;Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;Ljava/lang/String;)V + public fun (Lsk/ainet/compile/hlo/StableHloOperationRegistry;Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/compile/hlo/MlirValidator;Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;Ljava/lang/String;Lsk/ainet/compile/target/OpGranularityPolicy;)V + public synthetic fun (Lsk/ainet/compile/hlo/StableHloOperationRegistry;Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/compile/hlo/MlirValidator;Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;Ljava/lang/String;Lsk/ainet/compile/target/OpGranularityPolicy;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun convert (Lsk/ainet/lang/graph/ComputeGraph;Ljava/lang/String;)Lsk/ainet/compile/hlo/StableHloModule; public static synthetic fun convert$default (Lsk/ainet/compile/hlo/StableHloConverter;Lsk/ainet/lang/graph/ComputeGraph;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/compile/hlo/StableHloModule; public final fun convertWithOptimization (Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/compile/hlo/StableHloOptimizer;)Lsk/ainet/compile/hlo/StableHloModule; @@ -190,7 +196,9 @@ public final class sk/ainet/compile/hlo/StableHloConverterFactory { public static final field INSTANCE Lsk/ainet/compile/hlo/StableHloConverterFactory; public static final fun createBasic ()Lsk/ainet/compile/hlo/StableHloConverter; public static final fun createBasic (Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;)Lsk/ainet/compile/hlo/StableHloConverter; - public static synthetic fun createBasic$default (Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;ILjava/lang/Object;)Lsk/ainet/compile/hlo/StableHloConverter; + public static final fun createBasic (Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;Ljava/lang/String;)Lsk/ainet/compile/hlo/StableHloConverter; + public static final fun createBasic (Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;Ljava/lang/String;Lsk/ainet/compile/target/OpGranularityPolicy;)Lsk/ainet/compile/hlo/StableHloConverter; + public static synthetic fun createBasic$default (Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;Ljava/lang/String;Lsk/ainet/compile/target/OpGranularityPolicy;ILjava/lang/Object;)Lsk/ainet/compile/hlo/StableHloConverter; public static final fun createCustom (Lsk/ainet/compile/hlo/StableHloOperationRegistry;)Lsk/ainet/compile/hlo/StableHloConverter; public static final fun createCustom (Lsk/ainet/compile/hlo/StableHloOperationRegistry;Lsk/ainet/compile/hlo/TypeMapper;)Lsk/ainet/compile/hlo/StableHloConverter; public static final fun createCustom (Lsk/ainet/compile/hlo/StableHloOperationRegistry;Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/compile/hlo/MlirValidator;)Lsk/ainet/compile/hlo/StableHloConverter; @@ -315,6 +323,7 @@ public final class sk/ainet/compile/hlo/TypeMapper { public final fun mapDType (Ljava/lang/String;)Ljava/lang/String; public final fun mapFunctionSignature (Ljava/util/List;Ljava/util/List;)Ljava/lang/String; public final fun mapTensorType (Lsk/ainet/lang/tensor/ops/TensorSpec;)Ljava/lang/String; + public final fun negInfBits (Ljava/lang/String;)Ljava/lang/String; } public final class sk/ainet/compile/hlo/converters/ActivationOperationsConverter : sk/ainet/compile/hlo/StableHloOperationConverter { @@ -323,6 +332,12 @@ public final class sk/ainet/compile/hlo/converters/ActivationOperationsConverter public fun getSupportedOperations ()Ljava/util/Set; } +public final class sk/ainet/compile/hlo/converters/ArgMaxOperationsConverter : sk/ainet/compile/hlo/StableHloOperationConverter { + public fun ()V + public fun convert (Lsk/ainet/lang/graph/GraphNode;Ljava/util/List;Lsk/ainet/compile/hlo/ConversionContext;)Lsk/ainet/compile/hlo/ConversionResult; + public fun getSupportedOperations ()Ljava/util/Set; +} + public final class sk/ainet/compile/hlo/converters/AttentionOperationsConverter : sk/ainet/compile/hlo/StableHloOperationConverter { public fun ()V public fun convert (Lsk/ainet/lang/graph/GraphNode;Ljava/util/List;Lsk/ainet/compile/hlo/ConversionContext;)Lsk/ainet/compile/hlo/ConversionResult; diff --git a/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api b/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api index 331ca256..82e87835 100644 --- a/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api +++ b/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api @@ -1,3 +1,12 @@ +public final class sk/ainet/compile/opt/CompilePhase : java/lang/Enum { + public static final field DAG Lsk/ainet/compile/opt/CompilePhase; + public static final field STABLE_HLO Lsk/ainet/compile/opt/CompilePhase; + public static final field TAPE Lsk/ainet/compile/opt/CompilePhase; + public static fun getEntries ()Lkotlin/enums/EnumEntries; + public static fun valueOf (Ljava/lang/String;)Lsk/ainet/compile/opt/CompilePhase; + public static fun values ()[Lsk/ainet/compile/opt/CompilePhase; +} + public abstract interface class sk/ainet/compile/opt/GraphOptimizationPass { public abstract fun apply (Lsk/ainet/lang/graph/ComputeGraph;)Lsk/ainet/compile/opt/GraphOptimizationResult; public abstract fun getName ()Ljava/lang/String; @@ -51,6 +60,31 @@ public final class sk/ainet/compile/opt/GraphPruningKt { public static final fun prunedToOutputs (Lsk/ainet/lang/graph/ComputeGraph;Ljava/util/Set;)Lsk/ainet/lang/graph/ComputeGraph; } +public final class sk/ainet/compile/opt/TargetOptimizationKt { + public static final fun dagPipelineFor (Ljava/lang/String;Ljava/util/List;I)Lsk/ainet/compile/opt/GraphOptimizationPipeline; + public static synthetic fun dagPipelineFor$default (Ljava/lang/String;Ljava/util/List;IILjava/lang/Object;)Lsk/ainet/compile/opt/GraphOptimizationPipeline; +} + +public abstract interface class sk/ainet/compile/opt/TargetOptimizer { + public fun dagPasses ()Ljava/util/List; + public abstract fun getTarget ()Ljava/lang/String; + public fun granularity ()Lsk/ainet/compile/target/OpGranularityPolicy; +} + +public final class sk/ainet/compile/opt/TargetOptimizer$DefaultImpls { + public static fun dagPasses (Lsk/ainet/compile/opt/TargetOptimizer;)Ljava/util/List; + public static fun granularity (Lsk/ainet/compile/opt/TargetOptimizer;)Lsk/ainet/compile/target/OpGranularityPolicy; +} + +public final class sk/ainet/compile/opt/TargetOptimizers { + public static final field INSTANCE Lsk/ainet/compile/opt/TargetOptimizers; + public final fun clear ()V + public final fun forTarget (Ljava/lang/String;)Ljava/util/List; + public final fun granularityFor (Ljava/lang/String;)Lsk/ainet/compile/target/OpGranularityPolicy; + public final fun register (Lsk/ainet/compile/opt/TargetOptimizer;)V + public final fun registerDagPasses (Ljava/lang/String;Lkotlin/jvm/functions/Function0;)V +} + public final class sk/ainet/compile/opt/passes/ConstantFoldingPass : sk/ainet/compile/opt/GraphOptimizationPass { public fun ()V public fun apply (Lsk/ainet/lang/graph/ComputeGraph;)Lsk/ainet/compile/opt/GraphOptimizationResult; @@ -79,6 +113,14 @@ public final class sk/ainet/compile/opt/passes/DtypeConstraintViolationException public fun (Ljava/lang/String;)V } +public final class sk/ainet/compile/opt/passes/DtypeForwardPropagationPass : sk/ainet/compile/opt/GraphOptimizationPass { + public fun ()V + public fun (Ljava/lang/String;)V + public synthetic fun (Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun apply (Lsk/ainet/lang/graph/ComputeGraph;)Lsk/ainet/compile/opt/GraphOptimizationResult; + public fun getName ()Ljava/lang/String; +} + public final class sk/ainet/compile/opt/passes/LLMFusionPass : sk/ainet/compile/opt/GraphOptimizationPass { public fun ()V public fun apply (Lsk/ainet/lang/graph/ComputeGraph;)Lsk/ainet/compile/opt/GraphOptimizationResult; diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index eba0f5f7..68950992 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -4102,6 +4102,7 @@ public final class sk/ainet/lang/tensor/ops/KspTensorOps : sk/ainet/lang/tensor/ public fun abs (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun add (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun addScalar (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Number;)Lsk/ainet/lang/tensor/Tensor; + public fun argMax (Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor; public fun avgPool2d (Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;Z)Lsk/ainet/lang/tensor/Tensor; public fun clamp (Lsk/ainet/lang/tensor/Tensor;FF)Lsk/ainet/lang/tensor/Tensor; public fun concat (Ljava/util/List;I)Lsk/ainet/lang/tensor/Tensor; @@ -4354,6 +4355,7 @@ public abstract interface class sk/ainet/lang/tensor/ops/TensorOps { public abstract fun abs (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public abstract fun add (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public abstract fun addScalar (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Number;)Lsk/ainet/lang/tensor/Tensor; + public abstract fun argMax (Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor; public abstract fun avgPool2d (Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;Z)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun avgPool2d$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;ZILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public abstract fun clamp (Lsk/ainet/lang/tensor/Tensor;FF)Lsk/ainet/lang/tensor/Tensor; @@ -4556,6 +4558,7 @@ public final class sk/ainet/lang/tensor/ops/VoidTensorOps : sk/ainet/lang/tensor public fun abs (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun add (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun addScalar (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Number;)Lsk/ainet/lang/tensor/Tensor; + public fun argMax (Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor; public fun avgPool2d (Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;Z)Lsk/ainet/lang/tensor/Tensor; public fun clamp (Lsk/ainet/lang/tensor/Tensor;FF)Lsk/ainet/lang/tensor/Tensor; public fun concat (Ljava/util/List;I)Lsk/ainet/lang/tensor/Tensor; diff --git a/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api b/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api index ab274f2d..b0ea7ea5 100644 --- a/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api +++ b/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api @@ -41,6 +41,8 @@ public final class sk/ainet/lang/dag/GraphDslOpsGraphDslKt { public static synthetic fun add$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; public static final fun addScalar (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/Number;Ljava/lang/String;)Lsk/ainet/lang/dag/GraphValue; public static synthetic fun addScalar$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/Number;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; + public static final fun argMax (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;ILjava/lang/String;)Lsk/ainet/lang/dag/GraphValue; + public static synthetic fun argMax$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;ILjava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; public static final fun avgPool2d (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;ZLjava/lang/String;)Lsk/ainet/lang/dag/GraphValue; public static synthetic fun avgPool2d$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;ZLjava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; public static final fun clamp (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;FFLjava/lang/String;)Lsk/ainet/lang/dag/GraphValue; From c9e04452b6ad1a5dafa897ed6282d0442662c0e7 Mon Sep 17 00:00:00 2001 From: michal harakal Date: Wed, 8 Jul 2026 16:50:17 +0200 Subject: [PATCH 3/3] fix(ops): argMax eager result must be a portable float tensor (native/wasm) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The eager DefaultCpuOps.argMax returned an Int32 payload inside a tensor. JVM tolerated the boxing, but Kotlin/Native (linuxX64) and Wasm threw `ClassCastException: Int cannot be cast to Float` when reading .data — a real consumer bug, not just a test issue (surfaced by `allTests` in CI). Fix: the eager kernel now materializes the indices as index-valued floats in the input dtype, so the result is a consistent Tensor readable on every target. The traced/compiled form is unchanged — still a real i32 tensor (VoidTensorOps + ArgMaxOperationsConverter), so the board vmfb path is unaffected. Verified: DefaultCpuOpsArgMaxTest green on linuxX64 + jvm. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 11 +++++++---- .../ainet/exec/tensor/ops/DefaultCpuOpsArgMaxTest.kt | 7 +++---- .../kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt | 8 ++++++-- 3 files changed, 16 insertions(+), 10 deletions(-) diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt index 7070b21e..5bdd0d04 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt @@ -2079,12 +2079,15 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory var d = reduced.size - 1 while (d >= 0) { outIdx[d]++; if (outIdx[d] < reduced[d]) break; outIdx[d] = 0; d-- } } + // Eager result: store the indices as index-valued floats in the INPUT dtype so the tensor is + // a consistent Tensor readable on every target (an Int32 payload inside a + // tensor throws ClassCastException on Kotlin/Native + Wasm). The traced/compiled form is a + // real i32 tensor — see VoidTensorOps.argMax + ArgMaxOperationsConverter (emits stablehlo i32). + val floats = FloatArray(outCount) { indices[it].toFloat() } @Suppress("UNCHECKED_CAST") - val int32 = sk.ainet.lang.types.Int32::class as KClass - @Suppress("UNCHECKED_CAST") - val outData = dataFactory.fromIntArray(outShape, int32, indices) + val outData = dataFactory.fromFloatArray(outShape, tensor.dtype, floats) as sk.ainet.lang.tensor.data.TensorData - return CpuTensor(outData, this, int32, GradState(requiresGrad = false)) + return CpuTensor(outData, this, tensor.dtype, GradState(requiresGrad = false)) } @TensorOp() diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsArgMaxTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsArgMaxTest.kt index fc8b48e3..068ac02e 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsArgMaxTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsArgMaxTest.kt @@ -8,14 +8,13 @@ import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.dsl.tensor import sk.ainet.lang.types.FP32 -import sk.ainet.lang.types.Int32 class DefaultCpuOpsArgMaxTest { private val ctx = DirectCpuExecutionContext() private val ops get() = ctx.ops - // argMax returns an Int32-backed tensor (typed by the op signature); - // read the raw index via the underlying TensorData. + // Eager argMax returns index-valued floats in the input dtype (portable across JVM/native/wasm); + // read them back as ints. private fun idx(t: Tensor, vararg i: Int): Int = (t.data.get(*i) as Number).toInt() @Test @@ -26,7 +25,7 @@ class DefaultCpuOpsArgMaxTest { val t = tensor { shape(2, 3) { init { v[it[0] * 3 + it[1]] } } } val r = ops.argMax(t, dim = -1) assertEquals(Shape(2), r.shape) - assertEquals(Int32::class, r.dtype) // runtime dtype is Int32 despite the static type + assertEquals(FP32::class, r.dtype) // eager result is the input dtype (index-valued floats) assertEquals(2, idx(r, 0)) assertEquals(0, idx(r, 1)) } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt index 7f54239d..fe614093 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt @@ -196,8 +196,12 @@ public interface TensorOps { /** * Index of the maximum value along [dim]. The reduced dimension is removed from * the output (no keepdim); ties resolve to the LOWEST index (numpy/greedy - * semantics). The output holds `Int32` indices, so downstream consumers should - * treat it as an index tensor. Non-differentiable by design (index selection). + * semantics). Non-differentiable by design (index selection). + * + * Indices are integers: the traced/compiled form is a real `i32` tensor + * (see the StableHLO lowering), while the eager CPU path materializes them as + * index-valued floats in the input dtype so the result is a portable + * `Tensor` (an i32 payload in a float tensor is unreadable on native/wasm). * * Like [scaledDotProductAttention], this stays a single op and is lowered at the * StableHLO stage (`ArgMaxOperationsConverter`: reduce-max → broadcast → compare →