diff --git a/CHANGELOG.md b/CHANGELOG.md index 5931c9a9c..b641c1113 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,16 @@ ## [Unreleased] +### Added +- **`TensorOps.erf` and exact GELU** (#1342): `gelu(tensor, approximate: Boolean = true)` keeps the tanh form as the default and selects the exact `0.5·x·(1 + erf(x/√2))` with `approximate = false`, the form BERT-class encoders and the Laya decision models were trained with (the two differ by up to 4.7e-4 near |x| ≈ 2.7). CPU: Abramowitz–Stegun 7.1.26 erf evaluated in double (|err| < 1.5e-7) with a JVM Vector API lane kernel; `GELU(approximate)` module, `network {}` `gelu(approximate)`, Java `TensorJavaOps.gelu(a, approximate)` / `erf(a)` and `SequentialModelBuilder.gelu(approximate)`; exact backward rule and `erfBackward`. +- **GELU carries its form as op metadata**: the flag is recorded as `parameters["approximate"]` on the graph node (`GeluOperation`, `ErfOperation`), read by `ComputeGraphExecutor`, the backward rule and the StableHLO converter, so the compiled path cannot silently diverge from eager execution. StableHLO lowers `erf` to `chlo.erf` and exact GELU to `0.5·x·(1 + chlo.erf(x·√½))`; one `GeluLowering` builder emits both forms. + +### Changed +- `TensorOps.gelu` gained the `approximate` parameter: implementors of the interface must update the override (callers are source-compatible through the default). The legacy recording decorator now records `gelu`/`erf` as graph operations instead of delegating them unrecorded. + +### Fixed +- `DefaultCpuOps.gelu` printed a literal `${tensor.dtype}` in its unsupported-dtype message. + ## [0.58.0] - 2026-10-12 Headline: **input front ends and model identity (SKEEP-006).** Tokenizers leave `skainet-io-core` diff --git a/docs/modules/ROOT/pages/reference/operators/generated/index.adoc b/docs/modules/ROOT/pages/reference/operators/generated/index.adoc index 292f4ea63..c2a05e8ff 100644 --- a/docs/modules/ROOT/pages/reference/operators/generated/index.adoc +++ b/docs/modules/ROOT/pages/reference/operators/generated/index.adoc @@ -1,6 +1,6 @@ = AI-NET Operators Reference -Generated from version `0.58.0` on 2026-10-12 +Generated from version `0.58.0` on 2026-10-14 == Operators by Modality diff --git a/docs/modules/ROOT/pages/reference/operators/generated/tensorops.adoc b/docs/modules/ROOT/pages/reference/operators/generated/tensorops.adoc index a9e088dce..2f5241233 100644 --- a/docs/modules/ROOT/pages/reference/operators/generated/tensorops.adoc +++ b/docs/modules/ROOT/pages/reference/operators/generated/tensorops.adoc @@ -778,12 +778,13 @@ fun silu(tensor:Tensor): Tensor [source,kotlin] ---- -fun gelu(tensor:Tensor): Tensor +fun gelu(tensor:Tensor, approximate:Boolean): Tensor ---- === Parameters * `tensor: Tensor` +* `approximate: Boolean` === Return Type @@ -1287,6 +1288,25 @@ fun expm1(tensor:Tensor): Tensor `Tensor` +== erf + +[.darc-none]#✖ Generated facts only (no human prose)# + +=== Signature + +[source,kotlin] +---- +fun erf(tensor:Tensor): Tensor +---- + +=== Parameters + +* `tensor: Tensor` + +=== Return Type + +`Tensor` + == sin [.darc-none]#✖ Generated facts only (no human prose)# diff --git a/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc b/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc index 8154df4ae..8ec87a6ca 100644 --- a/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc +++ b/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc @@ -1,7 +1,7 @@ = Operator Coverage Matrix :description: Cross-backend status for every operator function in SKaiNET. -Generated from `operators.json` version `0.58.0` on 2026-10-12. +Generated from `operators.json` version `0.58.0` on 2026-10-14. Rows are `Operator.function` pairs. The `Validated` column shows whether the function's documentation has been DARC-validated by a reviewer (see xref:contributing/darc-workflow.adoc[DARC workflow]). Remaining columns are backends that appear in any function's `statusByBackend` map — a missing entry means the backend makes no claim about the function (treat it as "unknown", not "not supported"). @@ -70,12 +70,13 @@ Rows are `Operator.function` pairs. The `Validated` column shows whether the fun | `TensorOps.indexSelect` | ✖ | — | — | — | `TensorOps.exp` | ✖ | — | — | — | `TensorOps.expm1` | ✖ | — | — | — +| `TensorOps.erf` | ✖ | — | — | — | `TensorOps.sin` | ✖ | — | — | — | `TensorOps.cos` | ✖ | — | — | — | `TensorOps.scaledDotProductAttention` | ✖ | — | — | — | `Similarity.cosineDistance` | ✖ | ✅ | ✅ | ✅ -| *Done* | *1 / 65* | *1 / 65* | *1 / 65* | *1 / 65* +| *Done* | *1 / 66* | *1 / 66* | *1 / 66* | *1 / 66* |=== Per-function detail including notes lives in xref:reference/operators/generated/index.adoc[Operator reference]. 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 cac2c8479..0ea2ba696 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 @@ -300,13 +300,14 @@ public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/context/sched protected final fun elementwise (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lkotlin/jvm/functions/Function3;)Lsk/ainet/lang/tensor/Tensor; public fun elu (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; protected fun ensureKernelProviders ()V + public fun erf (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun exp (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun expm1 (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun flatten (Lsk/ainet/lang/tensor/Tensor;II)Lsk/ainet/lang/tensor/Tensor; protected final fun floatResult (Lsk/ainet/lang/tensor/Shape;Lkotlin/reflect/KClass;[F[Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun gather (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor; public fun ge (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; - public fun gelu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; + public fun gelu (Lsk/ainet/lang/tensor/Tensor;Z)Lsk/ainet/lang/tensor/Tensor; protected final fun getDataFactory ()Lsk/ainet/lang/tensor/data/TensorDataFactory; public final fun getSchedule ()Lsk/ainet/context/schedule/Schedule; protected final fun gradStateFrom ([Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/GradState; 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 c56964d0e..6ff2f1625 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 @@ -2596,8 +2596,8 @@ public open class DefaultCpuOpsBase( } @TensorOp() - @InProgress("cpu", owner = "team:cpu", issue = "task-ops.md#op-gelu") - override fun gelu(tensor: Tensor): Tensor { + override fun gelu(tensor: Tensor, approximate: Boolean): Tensor { + if (!approximate) return geluExact(tensor) floatUnaryFast(tensor) { x -> val inner = x + 0.044715f * (x * x * x) 0.5f * x * (1.0f + kotlin.math.tanh(0.7978845608f * inner)) @@ -2615,7 +2615,39 @@ public open class DefaultCpuOpsBase( @Suppress("UNCHECKED_CAST") y as V } - else -> throw IllegalArgumentException("Unsupported dtype for gelu: ${'$'}{tensor.dtype}") + else -> throw IllegalArgumentException("Unsupported dtype for gelu: ${tensor.dtype}") + } + } + return newTensor(outData, tensor.dtype, tensor) + } + + /** Exact GELU `0.5·x·(1 + erf(x/√2))`; erf from [ErfKernel] (|err| < 1.5e-7, evaluated in double). */ + private fun geluExact(tensor: Tensor): Tensor { + floatUnaryFast(tensor) { x -> ErfKernel.geluExact(x) }?.let { return it } + val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> + when (tensor.dtype) { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + val x = tensor.data.get(*idx) as Float + @Suppress("UNCHECKED_CAST") + ErfKernel.geluExact(x) as V + } + else -> throw IllegalArgumentException("Unsupported dtype for gelu: ${tensor.dtype}") + } + } + return newTensor(outData, tensor.dtype, tensor) + } + + @TensorOp() + override fun erf(tensor: Tensor): Tensor { + floatUnaryFast(tensor) { x -> ErfKernel.erf(x.toDouble()).toFloat() }?.let { return it } + val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> + when (tensor.dtype) { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + val x = tensor.data.get(*idx) as Float + @Suppress("UNCHECKED_CAST") + ErfKernel.erf(x.toDouble()).toFloat() as V + } + else -> throw IllegalArgumentException("Unsupported dtype for erf: ${tensor.dtype}") } } return newTensor(outData, tensor.dtype, tensor) diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/ErfKernel.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/ErfKernel.kt new file mode 100644 index 000000000..aa19f9a81 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/ErfKernel.kt @@ -0,0 +1,67 @@ +package sk.ainet.exec.tensor.ops + +import kotlin.math.abs +import kotlin.math.exp +import kotlin.math.sqrt + +/** + * Scalar Gauss error function and the exact GELU built on it (#1342). + * + * `erf` is Abramowitz–Stegun 7.1.26 (|err| < 1.5e-7) evaluated in double precision; the float + * result absorbs that error below its last bit. The expression order is the one the consumers + * validated against their reference goldens, so the kernel reproduces them bit-for-bit. + * + * The `*Float` variants evaluate the same polynomial in single precision with the same operation + * order; they serve the scalar tail of the JVM Vector API lane kernel, so lanes and tail agree. + */ +internal object ErfKernel { + private const val P = 0.3275911 + private const val A1 = 0.254829592 + private const val A2 = -0.284496736 + private const val A3 = 1.421413741 + private const val A4 = -1.453152027 + private const val A5 = 1.061405429 + + const val P_F: Float = 0.3275911f + const val A1_F: Float = 0.254829592f + const val A2_F: Float = -0.284496736f + const val A3_F: Float = 1.421413741f + const val A4_F: Float = -1.453152027f + const val A5_F: Float = 1.061405429f + /** 1/√2 in single precision. */ + const val SQRT1_2_F: Float = 0.70710678f + + /** `erf(x)` with absolute error below 1.5e-7, odd, saturating at ±1. */ + fun erf(x: Double): Double { + // The polynomial's coefficients sum to 1 only to ~1e-9, so pin the origin (and keep the + // sign of zero) instead of returning that residue with the sign of +0. + if (x == 0.0) return x + val sign = if (x < 0) -1.0 else 1.0 + val a = abs(x) + val t = 1.0 / (1.0 + P * a) + val y = 1.0 - (((((A5 * t + A4) * t) + A3) * t + A2) * t + A1) * t * exp(-a * a) + return sign * y + } + + /** `0.5 · x · (1 + erf(x / √2))` in double precision, rounded once to float. */ + fun geluExact(x: Float): Float { + val xd = x.toDouble() + return (0.5 * xd * (1.0 + erf(xd / sqrt(2.0)))).toFloat() + } + + /** Single-precision erf with the lane kernel's operation order (see `DefaultCpuOpsJvm.erfVector`). */ + fun erfFloat(x: Float): Float { + if (x == 0f) return x + val a = abs(x) + val t = 1f / (a * P_F + 1f) + var poly = t * A5_F + A4_F + poly = poly * t + A3_F + poly = poly * t + A2_F + poly = poly * t + A1_F + val y = 1f - poly * t * exp(-(a * a)) + return if (x < 0f) -y else y + } + + /** Single-precision exact GELU matching the lane kernel. */ + fun geluExactFloat(x: Float): Float = 0.5f * x * (1f + erfFloat(x * SQRT1_2_F)) +} diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsActivationsTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsActivationsTest.kt index a4cb7b6df..76e83b577 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsActivationsTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsActivationsTest.kt @@ -151,4 +151,103 @@ class DefaultCpuOpsActivationsTest { } assertTrue(threw, "Expected IllegalArgumentException for Int32 tanh") } + + // ---- erf / exact GELU (#1342) ---------------------------------------- + + /** Abramowitz–Stegun 7.1.26 in double, the host reference the kernel must reproduce. */ + private fun erfRef(x: Double): Double { + val sign = if (x < 0) -1.0 else 1.0 + val a = kotlin.math.abs(x) + val t = 1.0 / (1.0 + 0.3275911 * a) + val poly = ((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592 + return sign * (1.0 - poly * t * kotlin.math.exp(-a * a)) + } + + private fun geluExactRef(x: Float): Float { + val xd = x.toDouble() + return (0.5 * xd * (1.0 + erfRef(xd / kotlin.math.sqrt(2.0)))).toFloat() + } + + private fun geluTanhRef(x: Float): Float { + val inner = x + 0.044715f * (x * x * x) + return 0.5f * x * (1.0f + kotlin.math.tanh(0.7978845608f * inner)) + } + + @Test + fun erf_matches_reference_table() { + val xs = floatArrayOf(0f, 0.5f, -0.5f, 1f, -1f, 1.5f, -1.5f, 2f, -2f, 3f, -3f, 4f, -4f, 6f, -6f) + val table = floatArrayOf( + 0f, 0.5204998778f, -0.5204998778f, 0.8427007929f, -0.8427007929f, 0.9661051465f, -0.9661051465f, + 0.9953222650f, -0.9953222650f, 0.9999779095f, -0.9999779095f, 0.9999999846f, -0.9999999846f, 1f, -1f, + ) + val out = cpuOps.erf(fTensor(Shape(xs.size), xs)) + assertEquals(Shape(xs.size), out.shape) + assertEquals(FP32::class, out.dtype) + for (i in xs.indices) { + assertAlmostEquals(table[i], out.data[i] as Float, 2e-7f, "erf(${xs[i]}) mismatch: ${out.data[i]}") + } + } + + @Test + fun erf_is_odd_and_saturates() { + val xs = FloatArray(41) { -4f + it * 0.2f } + val pos = cpuOps.erf(fTensor(Shape(xs.size), xs)) + val neg = cpuOps.erf(fTensor(Shape(xs.size), FloatArray(xs.size) { -xs[it] })) + for (i in xs.indices) { + assertAlmostEquals(-(pos.data[i] as Float), neg.data[i] as Float, 0f, "erf must be odd at x=${xs[i]}") + } + val sat = cpuOps.erf(fTensor(Shape(2), floatArrayOf(6f, -6f))) + assertEquals(1f, sat.data[0] as Float) + assertEquals(-1f, sat.data[1] as Float) + } + + @Test + fun gelu_exact_matches_host_formula() { + val xs = FloatArray(121) { -6f + it * 0.1f } + val out = cpuOps.gelu(fTensor(Shape(xs.size), xs), approximate = false) + for (i in xs.indices) { + // Bit-equal on the JVM; JS/Wasm evaluate Float in double precision, so allow one ulp there. + assertAlmostEquals(geluExactRef(xs[i]), out.data[i] as Float, 2e-7f, "exact gelu at x=${xs[i]}") + } + } + + @Test + fun gelu_default_is_unchanged_tanh() { + val xs = floatArrayOf(-3f, -2f, -1f, -0.5f, 0f, 0.5f, 1f, 2f, 3f) + val t = fTensor(Shape(xs.size), xs) + val byDefault = cpuOps.gelu(t) + val explicit = cpuOps.gelu(t, approximate = true) + for (i in xs.indices) { + assertEquals(byDefault.data[i] as Float, explicit.data[i] as Float, "default must equal approximate=true") + assertAlmostEquals(geluTanhRef(xs[i]), byDefault.data[i] as Float, 1e-6f, "tanh gelu at x=${xs[i]}") + } + } + + @Test + fun gelu_exact_vs_tanh_max_difference() { + val xs = FloatArray(1201) { -6f + it * 0.01f } + val t = fTensor(Shape(xs.size), xs) + val exact = cpuOps.gelu(t, approximate = false) + val tanh = cpuOps.gelu(t, approximate = true) + var maxDiff = 0f + var argMax = 0f + for (i in xs.indices) { + val d = kotlin.math.abs((exact.data[i] as Float) - (tanh.data[i] as Float)) + if (d > maxDiff) { maxDiff = d; argMax = xs[i] } + } + // The tanh approximation deviates from the exact form by ≈ 4.7e-4 near |x| ≈ 2.7. + assertTrue(maxDiff > 4.5e-4f && maxDiff < 5.0e-4f, "max |exact - tanh| = $maxDiff at x=$argMax") + assertTrue(kotlin.math.abs(argMax) in 2.5f..2.9f, "arg max at x=$argMax") + } + + @Test + fun erf_and_exact_gelu_unsupported_dtype_int32_throws() { + val input = iTensor(Shape(2), intArrayOf(1, 2)) + @Suppress("UNCHECKED_CAST") + val asInt = input as sk.ainet.lang.tensor.Tensor + var threw = 0 + try { cpuOps.erf(asInt) } catch (e: IllegalArgumentException) { threw++; assertTrue(e.message?.contains("Unsupported dtype for erf") == true) } + try { cpuOps.gelu(asInt, approximate = false) } catch (e: IllegalArgumentException) { threw++; assertTrue(e.message?.contains("Unsupported dtype for gelu: ") == true && e.message?.contains("Int32") == true, e.message) } + assertEquals(2, threw) + } } diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt index b07602f43..41341bfdb 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt @@ -859,6 +859,39 @@ internal class DefaultCpuOpsJvm( return floatResult(Shape(tensor.shape.dimensions.copyOf()), tensor.dtype, out) } + override fun erf(tensor: Tensor): Tensor { + vectorFloatUnary(tensor, { v -> erfVector(v) }, ErfKernel::erfFloat)?.let { return it } + return super.erf(tensor) + } + + override fun gelu(tensor: Tensor, approximate: Boolean): Tensor { + if (approximate) return super.gelu(tensor, true) + vectorFloatUnary(tensor, { v -> + val half = FloatVector.broadcast(floatSpecies, 0.5f) + val one = FloatVector.broadcast(floatSpecies, 1f) + half.mul(v).mul(one.add(erfVector(v.mul(ErfKernel.SQRT1_2_F)))) + }, ErfKernel::geluExactFloat)?.let { return it } + return super.gelu(tensor, false) + } + + /** + * Lane-wise Abramowitz–Stegun 7.1.26 erf in single precision — the same operation order as + * [ErfKernel.erfFloat], which serves the scalar tail, so lanes and tail agree up to the + * last-ulp difference between `VectorOperators.EXP` and `kotlin.math.exp`. + */ + private fun erfVector(x: FloatVector): FloatVector { + val one = FloatVector.broadcast(floatSpecies, 1f) + val a = x.abs() + val t = one.div(a.mul(ErfKernel.P_F).add(1f)) + var poly = t.mul(ErfKernel.A5_F).add(ErfKernel.A4_F) + poly = poly.mul(t).add(ErfKernel.A3_F) + poly = poly.mul(t).add(ErfKernel.A2_F) + poly = poly.mul(t).add(ErfKernel.A1_F) + val e = a.mul(a).neg().lanewise(VectorOperators.EXP) + val y = one.sub(poly.mul(t).mul(e)) + return y.blend(y.neg(), x.lt(0f)) + } + override fun sum(tensor: Tensor, dim: Int?): Tensor { if (dim == null) { vectorFloatReduceAllSum(tensor)?.let { return it } diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvmElementwiseTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvmElementwiseTest.kt index 670579e98..aed19bf1b 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvmElementwiseTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvmElementwiseTest.kt @@ -55,6 +55,24 @@ class DefaultCpuOpsJvmElementwiseTest { } } + @Test + fun erf_and_exact_gelu_match_scalar_reference_with_tails() { + val species = jdk.incubator.vector.FloatVector.SPECIES_PREFERRED.length() + val lengths = listOf(1, 3, 7, species - 1, species, species + 1, species * 2 + 5, 257) + val scalar = DefaultCpuOps(dataFactory) + val jvm = DefaultCpuOpsJvm(dataFactory) + for (n in lengths) { + val a = FloatArray(n) { -8f + 16f * it / n.coerceAtLeast(2) } + val tA = floatTensor(Shape(n), a) + // erf: lane-wise single-precision polynomial vs the double-evaluated scalar kernel + compareTensors(scalar.erf(tA), jvm.erf(tA), eps = 1e-6f) + // exact gelu: |x| up to 8, so allow two ulps at that magnitude + compareTensors(scalar.gelu(tA, approximate = false), jvm.gelu(tA, approximate = false), eps = 2e-6f) + // the tanh default must stay byte-identical: the JVM tier delegates to the base kernel + compareTensors(scalar.gelu(tA), jvm.gelu(tA), eps = 0f) + } + } + @Test fun vector_flag_on_and_off_produce_same_results() { val n = 257 // ensure tail 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 78eb53216..df1912047 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 @@ -381,6 +381,19 @@ internal class RecordingTensorOpsDecorator(private val base: TensorOps) : Tensor return out } + override fun gelu(tensor: Tensor, approximate: Boolean): Tensor { + val out = base.gelu(tensor, approximate) + // The flag is the node parameter every later stage reads (executor, backward, exporters). + record(GeluOperation(mapOf(GeluOperation.APPROXIMATE to approximate)), listOf(tensor), listOf(out)) + return out + } + + override fun erf(tensor: Tensor): Tensor { + val out = base.erf(tensor) + record(ErfOperation(), listOf(tensor), listOf(out)) + return out + } + // --- Shape ops --- override fun reshape(tensor: Tensor, newShape: Shape): Tensor { val out = base.reshape(tensor, newShape) @@ -440,7 +453,6 @@ internal class RecordingTensorOpsDecorator(private val base: TensorOps) : Tensor // Delegations for unrecorded/less critical ops override fun split(tensor: Tensor, splitSize: Int, dim: Int): List> = base.split(tensor, splitSize, dim) - override fun gelu(tensor: Tensor): Tensor = base.gelu(tensor) override fun leakyRelu(tensor: Tensor, negativeSlope: Float): Tensor = base.leakyRelu(tensor, negativeSlope) override fun elu(tensor: Tensor, alpha: Float): Tensor = base.elu(tensor, alpha) override fun avgPool2d(input: Tensor, kernelSize: Pair, stride: Pair, padding: Pair, countIncludePad: Boolean): Tensor = base.avgPool2d(input, kernelSize, stride, padding, countIncludePad) diff --git a/skainet-compile/skainet-compile-core/src/commonTest/kotlin/sk/ainet/tape/SimpleExecutionTapeTest.kt b/skainet-compile/skainet-compile-core/src/commonTest/kotlin/sk/ainet/tape/SimpleExecutionTapeTest.kt index dbc80c5bd..b9d311df5 100644 --- a/skainet-compile/skainet-compile-core/src/commonTest/kotlin/sk/ainet/tape/SimpleExecutionTapeTest.kt +++ b/skainet-compile/skainet-compile-core/src/commonTest/kotlin/sk/ainet/tape/SimpleExecutionTapeTest.kt @@ -12,6 +12,8 @@ import sk.ainet.lang.tensor.data.DenseTensorDataFactory import sk.ainet.lang.tensor.ops.AddOperation import sk.ainet.lang.tensor.ops.Conv1dOperation import sk.ainet.lang.tensor.ops.Conv3dOperation +import sk.ainet.lang.tensor.ops.ErfOperation +import sk.ainet.lang.tensor.ops.GeluOperation import sk.ainet.lang.tensor.ops.ReluOperation import sk.ainet.lang.tensor.ops.TensorOps import sk.ainet.lang.tensor.ops.VoidTensorOps @@ -62,6 +64,32 @@ class SimpleExecutionTapeTest { assertEquals("FP32", first.inputs[0].dtype) } + @Test + fun records_gelu_with_approximate_flag_and_erf() { + val shape = Shape(3) + val a = VoidOpsTensor(dataFactory.full(shape, FP32::class, 0.5f), FP32::class) + + val tape = Execution.withTape { + val ops = Execution.recordingOps(VoidTensorOps()) + val g = ops.gelu(a, approximate = false) + val h = ops.gelu(g) + ops.erf(h) + } + + assertEquals(3, tape.operations.size, "gelu, gelu, erf must all be recorded") + val exact = tape.operations[0].operation + assertTrue(exact is GeluOperation<*, *>, "First op should be GeluOperation") + assertEquals(false, exact.parameters[GeluOperation.APPROXIMATE], "exact form must carry approximate=false") + assertFalse((exact as GeluOperation<*, *>).approximate) + + val approx = tape.operations[1].operation + assertTrue(approx is GeluOperation<*, *>) + assertEquals(true, approx.parameters[GeluOperation.APPROXIMATE], "default must record approximate=true") + + assertTrue(tape.operations[2].operation is ErfOperation<*, *>, "Third op should be ErfOperation") + assertEquals("erf", tape.operations[2].operation.name) + } + @Test fun tape_copy_and_clear_behaviour() { val shape = Shape(1) 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 63f778a75..6c4474ebd 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 @@ -110,6 +110,7 @@ public final class sk/ainet/lang/graph/DefaultGradientTape : sk/ainet/lang/graph public fun divScalarBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public fun divideBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public fun eluBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; + public fun erfBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public fun expBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public fun expm1Backward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public fun flattenBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt index be52912e2..b7acd1356 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt +++ b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt @@ -914,6 +914,17 @@ public class DefaultGradientTape( } override fun geluBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> { + val approximate = attributes["approximate"] as? Boolean ?: true + if (!approximate) { + // Exact form: GELU(x) = x · Φ(x) with Φ(x) = 0.5 · (1 + erf(x / √2)) + // dGELU/dx = Φ(x) + x · φ(x), φ(x) = exp(−x² / 2) / √(2π) + val x = inputs[0] + val ops = x.ops + val cdf = ops.mulScalar(ops.addScalar(ops.erf(ops.mulScalar(x, 0.7071067811865476)), 1.0), 0.5) + val pdf = ops.mulScalar(ops.exp(ops.mulScalar(ops.multiply(x, x), -0.5)), 0.3989422804014327) + val gradX = ops.add(cdf, ops.multiply(x, pdf)) + return listOf(ops.multiply(upstream, gradX)) + } // GELU(x) = 0.5 * x * (1 + erf(x / sqrt(2))) // dGELU(x)/dx = 0.5 * (1 + erf(x / sqrt(2))) + (x / sqrt(2*pi)) * exp(-x^2 / 2) // Using approximation: @@ -1033,6 +1044,14 @@ public class DefaultGradientTape( return listOf(ops.multiply(upstream, expX)) } + override fun erfBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> { + // d(erf(x))/dx = (2 / √π) · exp(−x²) + val x = inputs[0] + val ops = x.ops + val gradX = ops.mulScalar(ops.exp(ops.mulScalar(ops.multiply(x, x), -1.0)), 1.1283791670955126) + return listOf(ops.multiply(upstream, gradX)) + } + override fun scaledDotProductAttentionBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> { // SDPA backward is complex; stub returns null gradients (not differentiable through this path). // Full backward would require recomputing attention weights from Q,K,V. @@ -1289,6 +1308,7 @@ public class DefaultGradientTape( "split" to ::splitBackward, "exp" to ::expBackward, "expm1" to ::expm1Backward, + "erf" to ::erfBackward, "sin" to ::sinBackward, "cos" to ::cosBackward, "tril" to ::trilBackward, diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt index 31d3253bd..b07902ece 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt +++ b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt @@ -159,7 +159,8 @@ public class ComputeGraphExecutor( // Unary math ops "relu" -> listOf(ops.relu(inputs[0])) "silu" -> listOf(ops.silu(inputs[0])) - "gelu" -> listOf(ops.gelu(inputs[0])) + "gelu" -> listOf(ops.gelu(inputs[0], params["approximate"] as? Boolean ?: true)) + "erf" -> listOf(ops.erf(inputs[0])) "sigmoid" -> listOf(ops.sigmoid(inputs[0])) "exp" -> listOf(ops.exp(inputs[0])) "expm1" -> listOf(ops.expm1(inputs[0])) diff --git a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/grad/DefaultExecutionTapeTest.kt b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/grad/DefaultExecutionTapeTest.kt index 8629c3a90..f4e10324c 100644 --- a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/grad/DefaultExecutionTapeTest.kt +++ b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/grad/DefaultExecutionTapeTest.kt @@ -69,6 +69,52 @@ class DefaultExecutionTapeTest { assertTrue(approxEqual(analytic, num, tol = 5e-2f), "Analytic: ${analytic.joinToString()}, Numeric: ${num.joinToString()}") } + @Test + fun gelu_exact_gradient_matches_finite_difference() { + val trainCtx = createTrainCtx() + val values = floatArrayOf(-2f, -0.5f, 0f, 0.5f, 2f) + val a = trainCtx.fromFloatArray(Shape(values.size), FP32::class, values).withRequiresGrad() + + val pair = trainCtx.record { a.gelu(approximate = false).sum() } + val tape = pair.first as DefaultGradientTape + tape.computeGradients(targets = listOf(pair.second), sources = listOf(a)) + val analytic = fbuf(a.grad!!) + + val eps = 1e-3f + val num = FloatArray(values.size) + for (i in values.indices) { + val plus = values.copyOf(); plus[i] += eps + val minus = values.copyOf(); minus[i] -= eps + val lPlus = fbuf(trainCtx.fromFloatArray(Shape(values.size), FP32::class, plus).gelu(approximate = false).sum())[0] + val lMinus = fbuf(trainCtx.fromFloatArray(Shape(values.size), FP32::class, minus).gelu(approximate = false).sum())[0] + num[i] = (lPlus - lMinus) / (2 * eps) + } + assertTrue(approxEqual(analytic, num, tol = 1e-2f), "Analytic: ${analytic.joinToString()}, Numeric: ${num.joinToString()}") + } + + @Test + fun erf_gradient_matches_finite_difference() { + val trainCtx = createTrainCtx() + val values = floatArrayOf(-2f, -0.5f, 0f, 0.5f, 2f) + val a = trainCtx.fromFloatArray(Shape(values.size), FP32::class, values).withRequiresGrad() + + val pair = trainCtx.record { a.erf().sum() } + val tape = pair.first as DefaultGradientTape + tape.computeGradients(targets = listOf(pair.second), sources = listOf(a)) + val analytic = fbuf(a.grad!!) + + val eps = 1e-3f + val num = FloatArray(values.size) + for (i in values.indices) { + val plus = values.copyOf(); plus[i] += eps + val minus = values.copyOf(); minus[i] -= eps + val lPlus = fbuf(trainCtx.fromFloatArray(Shape(values.size), FP32::class, plus).erf().sum())[0] + val lMinus = fbuf(trainCtx.fromFloatArray(Shape(values.size), FP32::class, minus).erf().sum())[0] + num[i] = (lPlus - lMinus) / (2 * eps) + } + assertTrue(approxEqual(analytic, num, tol = 1e-2f), "Analytic: ${analytic.joinToString()}, Numeric: ${num.joinToString()}") + } + @Test fun concat_gradient_matches_finite_difference() { val trainCtx = createTrainCtx() 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 b89978ef4..d938a9940 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 @@ -233,7 +233,8 @@ private class TestTensorOps : TensorOps { override fun sigmoid(tensor: Tensor): Tensor = tensor override fun tanh(tensor: Tensor): Tensor = tensor override fun silu(tensor: Tensor): Tensor = tensor - override fun gelu(tensor: Tensor): Tensor = tensor + override fun gelu(tensor: Tensor, approximate: Boolean): Tensor = tensor + override fun erf(tensor: Tensor): Tensor = tensor 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 diff --git a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/GeluExactRecordingTest.kt b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/GeluExactRecordingTest.kt new file mode 100644 index 000000000..2200b8daf --- /dev/null +++ b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/GeluExactRecordingTest.kt @@ -0,0 +1,88 @@ +package sk.ainet.exec + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue +import sk.ainet.exec.tensor.ops.DefaultCpuOps +import sk.ainet.lang.graph.DefaultExecutionTape +import sk.ainet.lang.graph.DefaultGraphExecutionContext +import sk.ainet.lang.graph.exec.ComputeGraphExecutor +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.VoidOpsTensor +import sk.ainet.lang.tensor.data.DenseTensorDataFactory +import sk.ainet.lang.types.FP32 + +/** + * #1342 acceptance: the GELU form chosen on the eager (DIRECT) path is recorded as the node + * parameter `approximate`, and replaying the graph (OPTIMIZED) reproduces the eager output. + */ +class GeluExactRecordingTest { + + private val dataFactory = DenseTensorDataFactory() + + private fun tensor(shape: Shape, seed: Int): VoidOpsTensor { + val values = FloatArray(shape.volume) { i -> (((i * 37 + seed * 11) % 29) - 14) / 7f } + return VoidOpsTensor(dataFactory.fromFloatArray(shape, FP32::class, values), FP32::class) + } + + private fun floats(t: Tensor<*, *>): FloatArray = t.data.copyToFloatArray() + + @Test + fun exact_gelu_is_recorded_with_its_flag_and_replays_bit_identical() { + val ctx = DefaultGraphExecutionContext.tape(DefaultCpuOps(dataFactory)) + val x = tensor(Shape(4, 32), 1) + val w = tensor(Shape(32, 24), 2) + val b = tensor(Shape(4, 24), 3) + val w2 = tensor(Shape(24, 8), 4) + + ctx.startRecording() + val ops = ctx.ops + val h = ops.add(ops.matmul(x, w), b) + val g = ops.gelu(h, approximate = false) + val eager = ops.matmul(g, w2) + val tape = ctx.stopRecording() as DefaultExecutionTape + + // The trace carries the flag … + val geluTrace = tape.traces.single { it.opType == "gelu" } + assertEquals(false, geluTrace.attributes["approximate"]) + + // … and so does the graph node built from it. + val graph = tape.toComputeGraph(synthesizeExternalInputs = true, embedConstants = false) + val geluNode = graph.nodes.single { it.operation.name == "gelu" } + assertEquals(false, geluNode.operation.parameters["approximate"]) + + // Replay through the graph executor on the same kernels: external inputs keyed by trace id. + val session = ctx.session + val inputs = mapOf( + session.refOf(x).id to x, session.refOf(w).id to w, + session.refOf(b).id to b, session.refOf(w2).id to w2, + ) + val outputs = ComputeGraphExecutor(graph, DefaultCpuOps(dataFactory)).execute(inputs) + val lastMatmul = graph.nodes.filter { it.operation.name == "matmul" }.maxBy { it.id.substringAfter('n').substringBefore('_').toInt() } + val replayed = floats(outputs.getValue(lastMatmul.id)) + val expected = floats(eager) + + assertEquals(expected.size, replayed.size) + var maxDiff = 0f + for (i in expected.indices) maxDiff = maxOf(maxDiff, kotlin.math.abs(expected[i] - replayed[i])) + assertTrue(maxDiff <= 1e-6f, "DIRECT vs OPTIMIZED max |Δ| = $maxDiff") + + // Sanity: the exact form really was used — the tanh form differs on this block. + val tanhOut = floats(DefaultCpuOps(dataFactory).gelu(h, approximate = true)) + val exactOut = floats(g) + var formDiff = 0f + for (i in exactOut.indices) formDiff = maxOf(formDiff, kotlin.math.abs(exactOut[i] - tanhOut[i])) + assertTrue(formDiff > 1e-5f, "exact and tanh GELU should differ on this block, max |Δ| = $formDiff") + } + + @Test + fun default_gelu_records_approximate_true() { + val ctx = DefaultGraphExecutionContext.tape(DefaultCpuOps(dataFactory)) + val x = tensor(Shape(2, 3), 5) + ctx.startRecording() + ctx.ops.gelu(x) + val tape = ctx.stopRecording() as DefaultExecutionTape + assertEquals(true, tape.traces.single { it.opType == "gelu" }.attributes["approximate"]) + } +} 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 e5508232e..12a5e408e 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 @@ -398,12 +398,32 @@ public final class sk/ainet/compile/hlo/converters/ConstantOperationsConverter : public fun getSupportedOperations ()Ljava/util/Set; } +public final class sk/ainet/compile/hlo/converters/ContextEmitter : sk/ainet/compile/hlo/converters/HloEmitter { + public fun (Lsk/ainet/compile/hlo/ConversionContext;)V + public final fun getLines ()Ljava/util/List; + public fun newValue ()Ljava/lang/String; + public fun op (Ljava/lang/String;)V +} + public final class sk/ainet/compile/hlo/converters/GatherOperationsConverter : 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/GeluLowering { + public static final field INSTANCE Lsk/ainet/compile/hlo/converters/GeluLowering; + public static final field SQRT1_2 Ljava/lang/String; + public final fun emit (Lsk/ainet/compile/hlo/converters/HloEmitter;Ljava/lang/String;Ljava/lang/String;Z)Ljava/lang/String; + public final fun emitExact (Lsk/ainet/compile/hlo/converters/HloEmitter;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String; + public final fun emitTanh (Lsk/ainet/compile/hlo/converters/HloEmitter;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String; +} + +public abstract interface class sk/ainet/compile/hlo/converters/HloEmitter { + public abstract fun newValue ()Ljava/lang/String; + public abstract fun op (Ljava/lang/String;)V +} + public final class sk/ainet/compile/hlo/converters/LegacyOperationsConverter : 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-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt index 192657f61..11a3a16cf 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt @@ -256,8 +256,8 @@ public class ActivationOperationsConverter : StableHloOperationConverter { } /** - * Convert GELU activation using approximation. - * GELU(x) ≈ 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x^3))) + * Convert GELU. `parameters["approximate"]` (default `true`) selects the tanh chain; `false` + * the exact `chlo.erf` form — both emitted by [GeluLowering], the single source of the formula. */ private fun convertGelu( node: GraphNode, @@ -270,49 +270,21 @@ public class ActivationOperationsConverter : StableHloOperationConverter { "Unsupported gelu arity for node ${node.id}" ) } - + val outputSpec = node.outputs.firstOrNull() - val outputType = outputSpec?.let { context.getTypeMapper().mapTensorType(it) } + val outputType = outputSpec?.let { context.getTypeMapper().mapTensorType(it) } ?: "tensor" - - val halfConst = context.nextTempValue() - val oneConst = context.nextTempValue() - val coeff1Const = context.nextTempValue() // sqrt(2/π) ≈ 0.7978845608 - val coeff2Const = context.nextTempValue() // 0.044715 - val xSquared = context.nextTempValue() - val xCubed = context.nextTempValue() - val cubicTerm = context.nextTempValue() - val innerSum = context.nextTempValue() - val scaledSum = context.nextTempValue() - val tanhValue = context.nextTempValue() - val onePlusTanh = context.nextTempValue() - val halfX = context.nextTempValue() - val resultValue = context.nextTempValue() - - val operations = listOf( - "$halfConst = stablehlo.constant dense<0.5> : $outputType", - "$oneConst = stablehlo.constant dense<1.0> : $outputType", - "$coeff1Const = stablehlo.constant dense<0.7978845608> : $outputType", - "$coeff2Const = stablehlo.constant dense<0.044715> : $outputType", - "$xSquared = stablehlo.multiply ${operands[0]}, ${operands[0]} : $outputType", - "$xCubed = stablehlo.multiply $xSquared, ${operands[0]} : $outputType", - "$cubicTerm = stablehlo.multiply $coeff2Const, $xCubed : $outputType", - "$innerSum = stablehlo.add ${operands[0]}, $cubicTerm : $outputType", - "$scaledSum = stablehlo.multiply $coeff1Const, $innerSum : $outputType", - "$tanhValue = stablehlo.tanh $scaledSum : $outputType", - "$onePlusTanh = stablehlo.add $oneConst, $tanhValue : $outputType", - "$halfX = stablehlo.multiply $halfConst, ${operands[0]} : $outputType", - "$resultValue = stablehlo.multiply $halfX, $onePlusTanh : $outputType" - ) - - operations.forEach { context.emitOperation(it) } - + val approximate = (node.operation.parameters["approximate"] as? Boolean) ?: true + + val emitter = ContextEmitter(context) + val resultValue = GeluLowering.emit(emitter, operands[0], outputType, approximate) + return ConversionResult.Success( outputValueName = resultValue, - emittedOperations = operations + emittedOperations = emitter.lines ) } - + /** * Convert Swish activation (also known as SiLU). * Swish(x) = x * sigmoid(x) diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/GeluLowering.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/GeluLowering.kt new file mode 100644 index 000000000..5317b3284 --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/GeluLowering.kt @@ -0,0 +1,103 @@ +package sk.ainet.compile.hlo.converters + +import sk.ainet.compile.hlo.ConversionContext + +/** + * Minimal emission surface a lowering builder writes through: fresh SSA names and operation + * lines. [ConversionContext] is one implementation; a buffer that collects the body of a + * `func.func private` decomposition is another, so one builder serves both the inline form and + * a `stablehlo.composite` decomposition without the two ever drifting apart. + */ +public interface HloEmitter { + /** A fresh, module-unique SSA value name. */ + public fun newValue(): String + + /** Emit one operation line. */ + public fun op(line: String) +} + +/** [HloEmitter] over a [ConversionContext]; remembers the emitted lines for the conversion result. */ +public class ContextEmitter(private val context: ConversionContext) : HloEmitter { + private val _lines = mutableListOf() + + /** Lines emitted through this emitter, in order. */ + public val lines: List get() = _lines + + override fun newValue(): String = context.nextTempValue() + + override fun op(line: String) { + _lines += line + context.emitOperation(line) + } +} + +/** + * The GELU lowering (#1342), one builder for both forms: + * + * - `approximate = true` (default): `0.5·x·(1 + tanh(√(2/π)·(x + 0.044715·x³)))` in StableHLO primitives. + * - `approximate = false`: `0.5·x·(1 + erf(x/√2))` with `chlo.erf`. + * + * Constants are full-shape splats of the output type, as the other activation lowerings emit. + */ +public object GeluLowering { + + /** 1/√2 as the converter prints it. */ + public const val SQRT1_2: String = "0.70710678118654752" + + public fun emit(e: HloEmitter, x: String, type: String, approximate: Boolean): String = + if (approximate) emitTanh(e, x, type) else emitExact(e, x, type) + + /** Tanh approximation: byte-for-byte the chain the converter emitted before `approximate` existed. */ + public fun emitTanh(e: HloEmitter, x: String, type: String): String { + val halfConst = e.newValue() + val oneConst = e.newValue() + val coeff1Const = e.newValue() // sqrt(2/π) ≈ 0.7978845608 + val coeff2Const = e.newValue() // 0.044715 + val xSquared = e.newValue() + val xCubed = e.newValue() + val cubicTerm = e.newValue() + val innerSum = e.newValue() + val scaledSum = e.newValue() + val tanhValue = e.newValue() + val onePlusTanh = e.newValue() + val halfX = e.newValue() + val resultValue = e.newValue() + + e.op("$halfConst = stablehlo.constant dense<0.5> : $type") + e.op("$oneConst = stablehlo.constant dense<1.0> : $type") + e.op("$coeff1Const = stablehlo.constant dense<0.7978845608> : $type") + e.op("$coeff2Const = stablehlo.constant dense<0.044715> : $type") + e.op("$xSquared = stablehlo.multiply $x, $x : $type") + e.op("$xCubed = stablehlo.multiply $xSquared, $x : $type") + e.op("$cubicTerm = stablehlo.multiply $coeff2Const, $xCubed : $type") + e.op("$innerSum = stablehlo.add $x, $cubicTerm : $type") + e.op("$scaledSum = stablehlo.multiply $coeff1Const, $innerSum : $type") + e.op("$tanhValue = stablehlo.tanh $scaledSum : $type") + e.op("$onePlusTanh = stablehlo.add $oneConst, $tanhValue : $type") + e.op("$halfX = stablehlo.multiply $halfConst, $x : $type") + e.op("$resultValue = stablehlo.multiply $halfX, $onePlusTanh : $type") + return resultValue + } + + /** Exact form through `chlo.erf`. */ + public fun emitExact(e: HloEmitter, x: String, type: String): String { + val halfConst = e.newValue() + val oneConst = e.newValue() + val sqrt12Const = e.newValue() + val scaled = e.newValue() + val erfValue = e.newValue() + val onePlusErf = e.newValue() + val halfX = e.newValue() + val resultValue = e.newValue() + + e.op("$halfConst = stablehlo.constant dense<0.5> : $type") + e.op("$oneConst = stablehlo.constant dense<1.0> : $type") + e.op("$sqrt12Const = stablehlo.constant dense<$SQRT1_2> : $type") + e.op("$scaled = stablehlo.multiply $x, $sqrt12Const : $type") + e.op("$erfValue = chlo.erf $scaled : $type -> $type") + e.op("$onePlusErf = stablehlo.add $oneConst, $erfValue : $type") + e.op("$halfX = stablehlo.multiply $halfConst, $x : $type") + e.op("$resultValue = stablehlo.multiply $halfX, $onePlusErf : $type") + return resultValue + } +} diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/UnaryMathConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/UnaryMathConverter.kt index b7938c20f..c71e7e773 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/UnaryMathConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/UnaryMathConverter.kt @@ -35,6 +35,9 @@ public class UnaryMathConverter : StableHloOperationConverter { "round" to "stablehlo.round_nearest_even", "cos" to "stablehlo.cosine", "sin" to "stablehlo.sine", + // StableHLO has no erf; CHLO's `chlo.erf` is what `chlo-legalize-to-stablehlo` expands + // (IREE's StableHLO input pipeline runs it). CHLO unary ops print `: inType -> outType`. + "erf" to "chlo.erf", ) override val supportedOperations: Set = opMap.keys @@ -62,7 +65,11 @@ public class UnaryMathConverter : StableHloOperationConverter { ?: "tensor" val resultValue = context.nextTempValue() - val op = "$resultValue = $hloOp ${operands[0]} : $outputType" + val op = if (hloOp.startsWith("chlo.")) { + "$resultValue = $hloOp ${operands[0]} : $outputType -> $outputType" + } else { + "$resultValue = $hloOp ${operands[0]} : $outputType" + } context.emitOperation(op) return ConversionResult.Success( diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ActivationOperationsConverterTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ActivationOperationsConverterTest.kt index a194d80b5..3d89e8af3 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ActivationOperationsConverterTest.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ActivationOperationsConverterTest.kt @@ -132,6 +132,37 @@ class ActivationOperationsConverterTest { assertTrue(module.content.contains("tensor<2x3xf32>")) } + @Test + fun testGeluApproximateTrueParameterMatchesDefault() { + val byDefault = StableHloConverterFactory.createExtended() + .convert(createActivationGraph("gelu"), "test_gelu").content + val explicit = StableHloConverterFactory.createExtended() + .convert(createActivationGraphWithParameters("gelu", mapOf("approximate" to true)), "test_gelu").content + assertEquals(byDefault, explicit) + } + + @Test + fun testGeluExactOperation() { + val graph = createActivationGraphWithParameters("gelu", mapOf("approximate" to false)) + val converter = StableHloConverterFactory.createExtended() + val module = converter.convert(graph, "test_gelu_exact") + + // 0.5·x·(1 + erf(x/√2)) through chlo.erf (CHLO unary assembly: `: in -> out`) + assertTrue(module.content.contains("stablehlo.constant dense<0.5>")) + assertTrue(module.content.contains("stablehlo.constant dense<0.70710678118654752>")) + assertTrue(module.content.contains("chlo.erf")) + assertTrue(module.content.contains(": tensor<2x3xf32> -> tensor<2x3xf32>")) + assertFalse(module.content.contains("stablehlo.tanh")) + assertFalse(module.content.contains("0.044715")) + } + + @Test + fun testErfOperationLowersToChlo() { + val graph = createActivationGraph("erf") + val module = StableHloConverterFactory.createExtended().convert(graph, "test_erf") + assertTrue(module.content.contains("chlo.erf %arg0 : tensor<2x3xf32> -> tensor<2x3xf32>"), module.content) + } + @Test fun testGeluOperation() { val graph = createActivationGraph("gelu") 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 6165e033a..86868f9af 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 @@ -692,6 +692,8 @@ public final class sk/ainet/java/SequentialModelBuilder { public final fun flatten (II)Lsk/ainet/java/SequentialModelBuilder; public static synthetic fun flatten$default (Lsk/ainet/java/SequentialModelBuilder;IIILjava/lang/Object;)Lsk/ainet/java/SequentialModelBuilder; public final fun gelu ()Lsk/ainet/java/SequentialModelBuilder; + public final fun gelu (Z)Lsk/ainet/java/SequentialModelBuilder; + public static synthetic fun gelu$default (Lsk/ainet/java/SequentialModelBuilder;ZILjava/lang/Object;)Lsk/ainet/java/SequentialModelBuilder; public final fun input (I)Lsk/ainet/java/SequentialModelBuilder; public final fun relu ()Lsk/ainet/java/SequentialModelBuilder; public final fun sigmoid ()Lsk/ainet/java/SequentialModelBuilder; @@ -712,12 +714,15 @@ public final class sk/ainet/java/TensorJavaOps { public static final fun elu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public static final fun elu (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun elu$default (Lsk/ainet/lang/tensor/Tensor;FILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; + public static final fun erf (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public static final fun flatten (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public static final fun flatten (Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor; public static final fun flatten (Lsk/ainet/lang/tensor/Tensor;II)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun flatten$default (Lsk/ainet/lang/tensor/Tensor;IIILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static final fun ge (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; public static final fun gelu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; + public static final fun gelu (Lsk/ainet/lang/tensor/Tensor;Z)Lsk/ainet/lang/tensor/Tensor; + public static synthetic fun gelu$default (Lsk/ainet/lang/tensor/Tensor;ZILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static final fun leakyRelu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public static final fun leakyRelu (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun leakyRelu$default (Lsk/ainet/lang/tensor/Tensor;FILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; @@ -3103,9 +3108,10 @@ public final class sk/ainet/lang/nn/activations/ELU : sk/ainet/lang/nn/Module { public final class sk/ainet/lang/nn/activations/GELU : sk/ainet/lang/nn/Module { public fun ()V - public fun (Ljava/lang/String;)V - public synthetic fun (Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Ljava/lang/String;Z)V + public synthetic fun (Ljava/lang/String;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V public fun forward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/Tensor; + public final fun getApproximate ()Z public fun getModules ()Ljava/util/List; public fun getName ()Ljava/lang/String; } @@ -3657,12 +3663,14 @@ public final class sk/ainet/lang/nn/dsl/StageImpl : sk/ainet/lang/nn/dsl/NeuralN public final class sk/ainet/lang/nn/dsl/TensorOpsNetworkDslKt { public static final fun elu (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;FLjava/lang/String;)V public static synthetic fun elu$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;FLjava/lang/String;ILjava/lang/Object;)V + public static final fun erf (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;)V + public static synthetic fun erf$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;ILjava/lang/Object;)V public static final fun exp (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;)V public static synthetic fun exp$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;ILjava/lang/Object;)V public static final fun expm1 (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;)V public static synthetic fun expm1$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;ILjava/lang/Object;)V - public static final fun gelu (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;)V - public static synthetic fun gelu$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;ILjava/lang/Object;)V + public static final fun gelu (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ZLjava/lang/String;)V + public static synthetic fun gelu$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ZLjava/lang/String;ILjava/lang/Object;)V public static final fun leakyRelu (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;FLjava/lang/String;)V public static synthetic fun leakyRelu$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;FLjava/lang/String;ILjava/lang/Object;)V public static final fun relu (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;)V @@ -4598,12 +4606,14 @@ public final class sk/ainet/lang/tensor/TensorExtensionsKt { public static final fun div (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public static final fun elu (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun elu$default (Lsk/ainet/lang/tensor/Tensor;FILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; + public static final fun erf (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public static final fun exp (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public static final fun expm1 (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public static final fun flatten (Lsk/ainet/lang/tensor/Tensor;II)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun flatten$default (Lsk/ainet/lang/tensor/Tensor;IIILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static final fun ge (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; - public static final fun gelu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; + public static final fun gelu (Lsk/ainet/lang/tensor/Tensor;Z)Lsk/ainet/lang/tensor/Tensor; + public static synthetic fun gelu$default (Lsk/ainet/lang/tensor/Tensor;ZILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static final fun leakyRelu (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun leakyRelu$default (Lsk/ainet/lang/tensor/Tensor;FILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static final fun log (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; @@ -6304,6 +6314,7 @@ public abstract interface class sk/ainet/lang/tensor/ops/DifferentiableTensorOps public abstract fun divScalarBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public abstract fun divideBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public abstract fun eluBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; + public abstract fun erfBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public abstract fun expBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public abstract fun expm1Backward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public abstract fun flattenBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; @@ -6364,6 +6375,16 @@ public final class sk/ainet/lang/tensor/ops/DivideOperation : sk/ainet/lang/tens public fun validateInputs (Ljava/util/List;)Lsk/ainet/lang/tensor/ops/ValidationResult; } +public final class sk/ainet/lang/tensor/ops/ErfOperation : sk/ainet/lang/tensor/ops/BaseOperation { + public fun ()V + public fun (Ljava/util/Map;)V + public synthetic fun (Ljava/util/Map;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun clone (Ljava/util/Map;)Lsk/ainet/lang/tensor/ops/Operation; + public fun execute (Ljava/util/List;)Ljava/util/List; + public fun inferOutputs (Ljava/util/List;)Ljava/util/List; + public fun validateInputs (Ljava/util/List;)Lsk/ainet/lang/tensor/ops/ValidationResult; +} + public final class sk/ainet/lang/tensor/ops/FlattenOperation : sk/ainet/lang/tensor/ops/BaseOperation { public fun ()V public fun (Ljava/util/Map;)V @@ -6374,6 +6395,22 @@ public final class sk/ainet/lang/tensor/ops/FlattenOperation : sk/ainet/lang/ten public fun validateInputs (Ljava/util/List;)Lsk/ainet/lang/tensor/ops/ValidationResult; } +public final class sk/ainet/lang/tensor/ops/GeluOperation : sk/ainet/lang/tensor/ops/BaseOperation { + public static final field APPROXIMATE Ljava/lang/String; + public static final field Companion Lsk/ainet/lang/tensor/ops/GeluOperation$Companion; + public fun ()V + public fun (Ljava/util/Map;)V + public synthetic fun (Ljava/util/Map;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun clone (Ljava/util/Map;)Lsk/ainet/lang/tensor/ops/Operation; + public fun execute (Ljava/util/List;)Ljava/util/List; + public final fun getApproximate ()Z + public fun inferOutputs (Ljava/util/List;)Ljava/util/List; + public fun validateInputs (Ljava/util/List;)Lsk/ainet/lang/tensor/ops/ValidationResult; +} + +public final class sk/ainet/lang/tensor/ops/GeluOperation$Companion { +} + public final class sk/ainet/lang/tensor/ops/GenericOperation : sk/ainet/lang/tensor/ops/BaseOperation { public fun (Ljava/lang/String;Ljava/util/Map;Ljava/lang/String;)V public synthetic fun (Ljava/lang/String;Ljava/util/Map;Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V @@ -6415,12 +6452,13 @@ public final class sk/ainet/lang/tensor/ops/KspTensorOps : sk/ainet/lang/tensor/ public fun divScalar (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Number;)Lsk/ainet/lang/tensor/Tensor; public fun divide (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun elu (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; + public fun erf (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun exp (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun expm1 (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun flatten (Lsk/ainet/lang/tensor/Tensor;II)Lsk/ainet/lang/tensor/Tensor; public fun gather (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor; public fun ge (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; - public fun gelu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; + public fun gelu (Lsk/ainet/lang/tensor/Tensor;Z)Lsk/ainet/lang/tensor/Tensor; public fun indexSelect (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor; public fun leakyRelu (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; public fun log (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; @@ -6681,6 +6719,7 @@ public abstract interface class sk/ainet/lang/tensor/ops/TensorOps { public abstract fun divide (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public abstract fun elu (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun elu$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;FILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; + public abstract fun erf (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public abstract fun exp (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public abstract fun expm1 (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public abstract fun flatten (Lsk/ainet/lang/tensor/Tensor;II)Lsk/ainet/lang/tensor/Tensor; @@ -6688,7 +6727,8 @@ public abstract interface class sk/ainet/lang/tensor/ops/TensorOps { public abstract fun gather (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun gather$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public abstract fun ge (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; - public abstract fun gelu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; + public abstract fun gelu (Lsk/ainet/lang/tensor/Tensor;Z)Lsk/ainet/lang/tensor/Tensor; + public static synthetic fun gelu$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;ZILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public abstract fun indexSelect (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun indexSelect$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public abstract fun leakyRelu (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; @@ -6754,6 +6794,7 @@ public final class sk/ainet/lang/tensor/ops/TensorOps$DefaultImpls { public static synthetic fun elu$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;FILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun flatten$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;IIILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun gather$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; + public static synthetic fun gelu$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;ZILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun indexSelect$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun leakyRelu$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;FILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun logSoftmax$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; @@ -6891,12 +6932,13 @@ public final class sk/ainet/lang/tensor/ops/VoidTensorOps : sk/ainet/lang/tensor public fun divScalar (Lsk/ainet/lang/tensor/Tensor;Ljava/lang/Number;)Lsk/ainet/lang/tensor/Tensor; public fun divide (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun elu (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; + public fun erf (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun exp (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun expm1 (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun flatten (Lsk/ainet/lang/tensor/Tensor;II)Lsk/ainet/lang/tensor/Tensor; public fun gather (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor; public fun ge (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; - public fun gelu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; + public fun gelu (Lsk/ainet/lang/tensor/Tensor;Z)Lsk/ainet/lang/tensor/Tensor; public fun indexSelect (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;I)Lsk/ainet/lang/tensor/Tensor; public fun leakyRelu (Lsk/ainet/lang/tensor/Tensor;F)Lsk/ainet/lang/tensor/Tensor; public fun log (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/activations/GELU.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/activations/GELU.kt index fe5909f5f..92ebaff4e 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/activations/GELU.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/activations/GELU.kt @@ -6,10 +6,17 @@ import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.gelu import sk.ainet.lang.types.DType -public class GELU(override val name: String = "GELU") : Module() { +/** + * GELU activation module. [approximate] = `true` (default) applies the tanh form; `false` + * applies the exact erf form — see [sk.ainet.lang.tensor.ops.TensorOps.gelu]. + */ +public class GELU( + override val name: String = "GELU", + public val approximate: Boolean = true, +) : Module() { override val modules: List> get() = emptyList() override fun forward(input: Tensor, ctx: ExecutionContext): Tensor = - sk.ainet.lang.nn.hooks.withForwardHooks(ctx, this, input) { with(input) { gelu() } } + sk.ainet.lang.nn.hooks.withForwardHooks(ctx, this, input) { with(input) { gelu(approximate) } } } \ No newline at end of file diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/TensorExtensions.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/TensorExtensions.kt index eb9011990..26fe678d0 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/TensorExtensions.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/TensorExtensions.kt @@ -89,9 +89,10 @@ public fun Tensor.elu(alpha: Float = 1.0f): Tensor = public fun Tensor.sigmoid(): Tensor = ops.sigmoid(this) public fun Tensor.tanh(): Tensor = ops.tanh(this) public fun Tensor.silu(): Tensor = ops.silu(this) -public fun Tensor.gelu(): Tensor = ops.gelu(this) +public fun Tensor.gelu(approximate: Boolean = true): Tensor = ops.gelu(this, approximate) public fun Tensor.exp(): Tensor = ops.exp(this) public fun Tensor.expm1(): Tensor = ops.expm1(this) +public fun Tensor.erf(): Tensor = ops.erf(this) public fun Tensor.softmax(dim: Int = -1): Tensor = ops.softmax(this, dim) public fun Tensor.logSoftmax(dim: Int = -1): Tensor = ops.logSoftmax(this, dim) public fun Tensor.sum(dim: Int? = null): Tensor = ops.sum(this, dim) diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOperations.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOperations.kt index 4285085f6..440e63cf7 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOperations.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOperations.kt @@ -751,6 +751,83 @@ public class ReluOperation( override fun clone(newParameters: Map): Operation = ReluOperation(newParameters) } +/** + * GELU activation node. `parameters["approximate"]` (default `true`) selects the tanh form; + * `false` the exact erf form. The flag is what the graph executor, the backward rule and + * the exporters read, so the compiled path cannot silently diverge from eager execution. + */ +public class GeluOperation( + parameters: Map = emptyMap() +) : BaseOperation("gelu", "activation", parameters) { + + /** `true` = tanh approximation, `false` = exact `0.5·x·(1 + erf(x/√2))`. */ + public val approximate: Boolean get() = parameters[APPROXIMATE] as? Boolean ?: true + + override fun execute(inputs: List>): List> { + require(inputs.size == 1) { "GELU operation requires exactly 1 input" } + throw UnsupportedOperationException("Direct execution not supported in graph mode") + } + + override fun validateInputs(inputs: List): ValidationResult { + val errors = mutableListOf() + if (inputs.size != 1) errors += "GELU operation requires exactly 1 input, got ${inputs.size}" + val flag = parameters[APPROXIMATE] + if (flag != null && flag !is Boolean) errors += "GELU parameter '$APPROXIMATE' must be a Boolean, got $flag" + return if (errors.isEmpty()) ValidationResult.Valid else ValidationResult.Invalid(errors) + } + + override fun inferOutputs(inputs: List): List { + require(inputs.size == 1) { "GELU operation requires exactly 1 input" } + return listOf( + TensorSpec( + name = "gelu_output", + shape = inputs[0].shape, + dtype = inputs[0].dtype, + requiresGrad = inputs[0].requiresGrad + ) + ) + } + + override fun clone(newParameters: Map): Operation = GeluOperation(newParameters) + + public companion object { + /** Parameter key carrying the approximation flag. */ + public const val APPROXIMATE: String = "approximate" + } +} + +/** Element-wise Gauss error function node. */ +public class ErfOperation( + parameters: Map = emptyMap() +) : BaseOperation("erf", "math", parameters) { + + override fun execute(inputs: List>): List> { + require(inputs.size == 1) { "erf operation requires exactly 1 input" } + throw UnsupportedOperationException("Direct execution not supported in graph mode") + } + + override fun validateInputs(inputs: List): ValidationResult { + if (inputs.size != 1) { + return ValidationResult.Invalid(listOf("erf operation requires exactly 1 input, got ${inputs.size}")) + } + return ValidationResult.Valid + } + + override fun inferOutputs(inputs: List): List { + require(inputs.size == 1) { "erf operation requires exactly 1 input" } + return listOf( + TensorSpec( + name = "erf_output", + shape = inputs[0].shape, + dtype = inputs[0].dtype, + requiresGrad = inputs[0].requiresGrad + ) + ) + } + + override fun clone(newParameters: Map): Operation = ErfOperation(newParameters) +} + public class SoftmaxOperation( parameters: Map = emptyMap() ) : BaseOperation("softmax", "activation", 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 bc0a552f9..75bcd0977 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 @@ -219,9 +219,18 @@ public interface TensorOps { @Diff @ActivationDsl public fun silu(tensor: Tensor): Tensor + /** + * Gaussian Error Linear Unit. + * + * `approximate = true` (default) is the tanh form `0.5·x·(1 + tanh(√(2/π)·(x + 0.044715·x³)))`, + * numerically unchanged from earlier releases. `approximate = false` is the exact form + * `0.5·x·(1 + erf(x/√2))` used by BERT-class encoders trained with `hidden_act = "gelu"`; + * the two differ by at most ≈ 4.7e-4 (near |x| ≈ 2.7). The flag is recorded as the + * operation parameter `approximate` so the compiled path and exporters see the same choice. + */ @Diff @ActivationDsl - public fun gelu(tensor: Tensor): Tensor + public fun gelu(tensor: Tensor, approximate: Boolean = true): Tensor // Reduction operations @Diff @@ -357,6 +366,11 @@ public interface TensorOps { @ActivationDsl public fun expm1(tensor: Tensor): Tensor + /** Element-wise Gauss error function `erf(x) = 2/√π · ∫₀ˣ e^(−t²) dt`, odd, saturating at ±1. */ + @Diff + @ActivationDsl + public fun erf(tensor: Tensor): Tensor + // Trigonometric operations @Diff public fun sin(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 1920ffce3..973a8c8b7 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 @@ -388,7 +388,7 @@ public class VoidTensorOps : TensorOps { return VoidOpsTensor(resultData, tensor.dtype) } - override fun gelu(tensor: Tensor): Tensor { + override fun gelu(tensor: Tensor, approximate: Boolean): Tensor { // GELU activation function preserves shape val resultData = dataFactory.zeros(tensor.shape, tensor.dtype) return VoidOpsTensor(resultData, tensor.dtype) @@ -540,6 +540,11 @@ public class VoidTensorOps : TensorOps { return VoidOpsTensor(resultData, tensor.dtype) } + override fun erf(tensor: Tensor): Tensor { + val resultData = dataFactory.zeros(tensor.shape, tensor.dtype) + return VoidOpsTensor(resultData, tensor.dtype) + } + override fun sin(tensor: Tensor): Tensor { val resultData = dataFactory.zeros(tensor.shape, tensor.dtype) return VoidOpsTensor(resultData, tensor.dtype) diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilderCoverageTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilderCoverageTest.kt index c39ac7b4e..d39e817fe 100644 --- a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilderCoverageTest.kt +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilderCoverageTest.kt @@ -462,6 +462,21 @@ class NetworkBuilderCoverageTest { assertNotNull(net) } + @Test + fun testGeneratedGeluActivationDslExactForm() { + // The KSP-generated gelu(approximate, id) must default to the tanh form and accept the exact one + val net = definition { + sequential { + input(8) + dense(4) { weights { ones() }; bias { zeros() } } + gelu(approximate = false, id = "gelu_exact") + dense(2) { weights { ones() }; bias { zeros() } } + gelu(id = "gelu_default") + } + } + assertNotNull(net) + } + @Test fun testMultipleGeneratedActivationMethods() { // Test using multiple KSP-generated activation methods in one network @@ -475,7 +490,7 @@ class NetworkBuilderCoverageTest { dense(8) { weights { ones() }; bias { zeros() } } elu(1.0f, "elu_1") dense(4) { weights { ones() }; bias { zeros() } } - gelu("gelu_1") + gelu(id = "gelu_1") dense(2) { weights { ones() }; bias { zeros() } } sigmoid("output") } diff --git a/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/java/SequentialModelBuilder.kt b/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/java/SequentialModelBuilder.kt index b9f90efe7..3e15a64be 100644 --- a/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/java/SequentialModelBuilder.kt +++ b/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/java/SequentialModelBuilder.kt @@ -88,9 +88,10 @@ public class SequentialModelBuilder @JvmOverloads constructor( return this } - /** Adds a GELU activation. */ - public fun gelu(): SequentialModelBuilder { - layers.add(ActivationsWrapperModule({ tensor -> tensor.ops.gelu(tensor) }, "gelu")) + /** Adds a GELU activation; `approximate = false` selects the exact erf form. */ + @JvmOverloads + public fun gelu(approximate: Boolean = true): SequentialModelBuilder { + layers.add(ActivationsWrapperModule({ tensor -> tensor.ops.gelu(tensor, approximate) }, "gelu")) return this } diff --git a/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/java/TensorJavaOps.kt b/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/java/TensorJavaOps.kt index 384b706e2..6e3c8fb4c 100644 --- a/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/java/TensorJavaOps.kt +++ b/skainet-lang/skainet-lang-core/src/jvmMain/kotlin/sk/ainet/java/TensorJavaOps.kt @@ -162,12 +162,21 @@ public object TensorJavaOps { return ta.ops.silu(ta) } - /** GELU activation (Gaussian Error Linear Unit). */ + /** GELU activation (Gaussian Error Linear Unit); `approximate = false` selects the exact erf form. */ @JvmStatic - public fun gelu(a: Tensor<*, *>): Tensor<*, *> { + @JvmOverloads + public fun gelu(a: Tensor<*, *>, approximate: Boolean = true): Tensor<*, *> { + @Suppress("UNCHECKED_CAST") + val ta = a as Tensor + return ta.ops.gelu(ta, approximate) + } + + /** Element-wise Gauss error function. */ + @JvmStatic + public fun erf(a: Tensor<*, *>): Tensor<*, *> { @Suppress("UNCHECKED_CAST") val ta = a as Tensor - return ta.ops.gelu(ta) + return ta.ops.erf(ta) } /** Softmax along the given dimension. */ 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 f21d8b6c1..3bd72d832 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 @@ -71,6 +71,8 @@ public final class sk/ainet/lang/dag/GraphDslOpsGraphDslKt { public static synthetic fun divide$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 elu (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;FLjava/lang/String;)Lsk/ainet/lang/dag/GraphValue; public static synthetic fun elu$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;FLjava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; + public static final fun erf (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;)Lsk/ainet/lang/dag/GraphValue; + public static synthetic fun erf$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; public static final fun exp (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;)Lsk/ainet/lang/dag/GraphValue; public static synthetic fun exp$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; public static final fun expm1 (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;)Lsk/ainet/lang/dag/GraphValue; @@ -81,8 +83,8 @@ public final class sk/ainet/lang/dag/GraphDslOpsGraphDslKt { public static synthetic fun gather$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;ILjava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; public static final fun ge (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;FLjava/lang/String;)Lsk/ainet/lang/dag/GraphValue; public static synthetic fun ge$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;FLjava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; - public static final fun gelu (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;)Lsk/ainet/lang/dag/GraphValue; - public static synthetic fun gelu$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; + public static final fun gelu (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;ZLjava/lang/String;)Lsk/ainet/lang/dag/GraphValue; + public static synthetic fun gelu$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;ZLjava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; public static final fun indexSelect (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;ILjava/lang/String;)Lsk/ainet/lang/dag/GraphValue; public static synthetic fun indexSelect$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;ILjava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; public static final fun leakyRelu (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;FLjava/lang/String;)Lsk/ainet/lang/dag/GraphValue; diff --git a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/GraphDslGenerator.kt b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/GraphDslGenerator.kt index dee40bdda..b6748061a 100644 --- a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/GraphDslGenerator.kt +++ b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/GraphDslGenerator.kt @@ -104,7 +104,9 @@ class GraphDslGenerator( } else if (param.type == "sk.ainet.lang.tensor.ops.UpsampleMode") { " = UpsampleMode.Nearest" } else if (param.type == "kotlin.Boolean") { + // Name-keyed defaults (KSP cannot read the declared expression); keep in sync with TensorOps. if (param.name == "alignCorners") " = false" + else if (param.name == "approximate") " = true" else "" } else "" } else "" diff --git a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/MethodAnalyzer.kt b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/MethodAnalyzer.kt index 9cfbc0689..59d9b1244 100644 --- a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/MethodAnalyzer.kt +++ b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/MethodAnalyzer.kt @@ -315,7 +315,7 @@ class MethodAnalyzer { } // Unary operations - method.name in listOf("relu", "sigmoid", "tanh", "exp", "log", "abs", "neg", "silu", "softmax", "logSoftmax", "sqrt") && + method.name in listOf("relu", "sigmoid", "tanh", "exp", "log", "abs", "neg", "silu", "softmax", "logSoftmax", "sqrt", "gelu", "erf") && tensorParams.size == 1 && method.returnType.isTensor -> { val param = tensorParams[0].name "unary($param as Tensor<*, *>, result as Tensor<*, *>)" diff --git a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/NetworkDslGenerator.kt b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/NetworkDslGenerator.kt index f705771e0..c878a625d 100644 --- a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/NetworkDslGenerator.kt +++ b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/NetworkDslGenerator.kt @@ -153,11 +153,14 @@ class NetworkDslGenerator( } private fun getDefaultValue(param: ParameterInfo): String? { - // Handle known default values for activation parameters + // Handle known default values for activation parameters. KSP does not expose the + // default-value expression of an interface parameter, so this table is keyed by + // parameter name and must mirror the declaration in TensorOps. return when { param.name == "negativeSlope" && param.type.contains("Float") -> "0.01f" param.name == "alpha" && param.type.contains("Float") -> "1.0f" param.name == "dim" && param.type.contains("Int") -> "-1" + param.name == "approximate" && param.type.contains("Boolean") -> "true" param.defaultValue != null -> param.defaultValue param.isOptional -> when { param.type.contains("Float") -> "0.0f" @@ -175,6 +178,7 @@ class NetworkDslGenerator( "negativeSlope" -> "Controls the angle of the negative slope in LeakyReLU" "alpha" -> "Scale for the negative factor in ELU" "dim" -> "Dimension along which to compute the operation" + "approximate" -> "true = tanh approximation (default), false = exact erf form" else -> "Parameter for the activation function" } } diff --git a/skainet-test/skainet-test-groundtruth/src/commonMain/kotlin/sk/ainet/test/groundtruth/GroundTruthTestCase.kt b/skainet-test/skainet-test-groundtruth/src/commonMain/kotlin/sk/ainet/test/groundtruth/GroundTruthTestCase.kt index a9267e508..3deb3e3c6 100644 --- a/skainet-test/skainet-test-groundtruth/src/commonMain/kotlin/sk/ainet/test/groundtruth/GroundTruthTestCase.kt +++ b/skainet-test/skainet-test-groundtruth/src/commonMain/kotlin/sk/ainet/test/groundtruth/GroundTruthTestCase.kt @@ -151,5 +151,8 @@ public data class OperationParams( val negativeSlope: Float? = null, /** Alpha for ELU */ - val alpha: Float? = null + val alpha: Float? = null, + + /** GELU form: `false` = exact erf, `true`/null = tanh approximation */ + val approximate: Boolean? = null ) diff --git a/skainet-test/skainet-test-groundtruth/src/commonMain/kotlin/sk/ainet/test/groundtruth/OperationExecutor.kt b/skainet-test/skainet-test-groundtruth/src/commonMain/kotlin/sk/ainet/test/groundtruth/OperationExecutor.kt index 3789ce6a4..5988edb49 100644 --- a/skainet-test/skainet-test-groundtruth/src/commonMain/kotlin/sk/ainet/test/groundtruth/OperationExecutor.kt +++ b/skainet-test/skainet-test-groundtruth/src/commonMain/kotlin/sk/ainet/test/groundtruth/OperationExecutor.kt @@ -129,7 +129,11 @@ public class OperationExecutor( } "gelu" -> { val a = getSingleInput(inputs) - ops.gelu(a) + ops.gelu(a, params.approximate ?: true) + } + "erf" -> { + val a = getSingleInput(inputs) + ops.erf(a) } "softmax" -> { val a = getSingleInput(inputs) @@ -313,6 +317,7 @@ public class OperationParamsBuilder { private var endDim: Int? = null private var negativeSlope: Float? = null private var alpha: Float? = null + private var approximate: Boolean? = null public fun stride(h: Int, w: Int): OperationParamsBuilder { stride = h to w @@ -367,6 +372,11 @@ public class OperationParamsBuilder { return this } + public fun approximate(value: Boolean): OperationParamsBuilder { + approximate = value + return this + } + public fun alpha(value: Float): OperationParamsBuilder { alpha = value return this @@ -382,7 +392,8 @@ public class OperationParamsBuilder { startDim = startDim, endDim = endDim, negativeSlope = negativeSlope, - alpha = alpha + alpha = alpha, + approximate = approximate ) }