diff --git a/docs/modules/ROOT/pages/reference/operators/generated/index.adoc b/docs/modules/ROOT/pages/reference/operators/generated/index.adoc index 441bdbe33..dd9000036 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.62.4` on 2026-10-25 +Generated from version `0.62.4` on 2026-10-30 == Operators by Modality diff --git a/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc b/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc index 8638d49da..2a434d8bc 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.62.4` on 2026-10-25. +Generated from `operators.json` version `0.62.4` on 2026-10-30. Rows are `Operator.function` pairs. A gap is marked in source with `@NotImplemented(reason = …, issue = …)` (or `@InProgress`) and shows up here as ❌ (or 🚧) linked to its tracking issue; the PR that closes the gap removes the annotation. See xref:contributing/issue-taxonomy.adoc[Issue taxonomy] to pick one up. @@ -125,9 +125,9 @@ Every concrete `sk.ainet.lang.nn` `Module` subclass. ❌ means the layer's `forw | `Dropout` | `sk.ainet.lang.nn.layers` | ✅ | — | `LossModule` | `sk.ainet.lang.nn.loss` | ✅ | — | `BatchNormalization` | `sk.ainet.lang.nn.normalization` | ✅ | — -| `GroupNormalization` | `sk.ainet.lang.nn.normalization` | ❌ | forward helpers (group reshape, mean, variance) are TODO(); eps and gamma/beta are not applied — https://github.com/SKaiNET-developers/SKaiNET/issues/1379[#1379] +| `GroupNormalization` | `sk.ainet.lang.nn.normalization` | ✅ | — | `LayerNormalization` | `sk.ainet.lang.nn.normalization` | ✅ | — | `MLP` | `sk.ainet.lang.nn.topology` | ✅ | — -| *Done* | | *30 / 32* | +| *Done* | | *31 / 32* | |=== diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/lang/nn/normalization/GroupNormalizationTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/lang/nn/normalization/GroupNormalizationTest.kt new file mode 100644 index 000000000..ea482ae6b --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/lang/nn/normalization/GroupNormalizationTest.kt @@ -0,0 +1,236 @@ +package sk.ainet.lang.nn.normalization + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.math.sqrt +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +class GroupNormalizationTest { + + private val ctx = DirectCpuExecutionContext() + + private fun assertAlmostEquals(expected: Float, actual: Float, eps: Float = 1e-5f, msg: String = "") { + assertTrue(abs(expected - actual) <= eps, msg.ifEmpty { "Expected $expected, got $actual (diff=${abs(expected - actual)})" }) + } + + private fun tensor(shape: Shape, data: FloatArray): Tensor = + ctx.fromFloatArray(shape, FP32::class, data) + + /** + * Reference GroupNorm with torch.nn.GroupNorm semantics, computed in Double: + * per (sample, group) biased mean/variance over C/G * spatial elements, + * then per-channel affine. Input is flat NCHW order. + */ + private fun referenceGroupNorm( + input: FloatArray, + n: Int, + c: Int, + spatial: Int, + groups: Int, + eps: Double = 1e-5, + gamma: FloatArray? = null, + beta: FloatArray? = null, + ): FloatArray { + val channelsPerGroup = c / groups + val groupSize = channelsPerGroup * spatial + val out = FloatArray(input.size) + for (b in 0 until n) { + for (g in 0 until groups) { + val indices = buildList { + for (ch in g * channelsPerGroup until (g + 1) * channelsPerGroup) { + for (s in 0 until spatial) add((b * c + ch) * spatial + s) + } + } + val mean = indices.sumOf { input[it].toDouble() } / groupSize + val variance = indices.sumOf { (input[it] - mean) * (input[it] - mean) } / groupSize + val denom = sqrt(variance + eps) + for (i in indices) { + val ch = (i / spatial) % c + var v = (input[i] - mean) / denom + if (gamma != null) v = v * gamma[ch] + (beta?.get(ch)?.toDouble() ?: 0.0) + out[i] = v.toFloat() + } + } + } + return out + } + + private fun assertFlatClose(expected: FloatArray, actual: Tensor, shape: Shape, eps: Float = 1e-5f) { + assertEquals(shape, actual.shape) + val dims = shape.dimensions + var flat = 0 + val idx = IntArray(dims.size) + while (flat < expected.size) { + var rem = flat + for (d in dims.indices.reversed()) { + idx[d] = rem % dims[d] + rem /= dims[d] + } + assertAlmostEquals(expected[flat], actual.data.get(*idx) as Float, eps, "Mismatch at flat index $flat") + flat++ + } + } + + @Test + fun groupNorm_2d_noAffine_normalizesPerGroup() { + // (N=1, C=4), G=2 -> groups [1,2] and [3,4], each mean +/- 0.5, var 0.25 + val data = floatArrayOf(1f, 2f, 3f, 4f) + val input = tensor(Shape(1, 4), data) + val gn = GroupNormalization(numGroups = 2, numChannels = 4, affine = false, name = "gn") + + val output = gn.forward(input, ctx) + + val expected = referenceGroupNorm(data, n = 1, c = 4, spatial = 1, groups = 2) + assertFlatClose(expected, output, Shape(1, 4)) + // Sanity against hand math: (x - m) / sqrt(0.25 + 1e-5) ~= +/- 0.99998 + assertAlmostEquals(-0.99998f, output.data[0, 0] as Float, 1e-4f) + assertAlmostEquals(0.99998f, output.data[0, 1] as Float, 1e-4f) + } + + @Test + fun groupNorm_4d_noAffine_matchesReference() { + // (N=2, C=4, H=2, W=2), G=2 — varied values, groups span 2 channels * 4 spatial + val data = FloatArray(32) { i -> (i * 7 % 13) - 6f + i * 0.25f } + val input = tensor(Shape(2, 4, 2, 2), data) + val gn = GroupNormalization(numGroups = 2, numChannels = 4, affine = false, name = "gn4d") + + val output = gn.forward(input, ctx) + + val expected = referenceGroupNorm(data, n = 2, c = 4, spatial = 4, groups = 2) + assertFlatClose(expected, output, Shape(2, 4, 2, 2)) + } + + @Test + fun groupNorm_4d_withAffine_appliesPerChannelGammaBeta() { + // Distinct gamma/beta per channel: catches wrong broadcast axis. + val data = FloatArray(16) { i -> i.toFloat() } + val input = tensor(Shape(1, 4, 2, 2), data) + val gammaArr = floatArrayOf(1f, 2f, 3f, 4f) + val betaArr = floatArrayOf(0.5f, -0.5f, 1f, -1f) + val gn = GroupNormalization( + numGroups = 2, + numChannels = 4, + affine = true, + name = "gn_affine", + initGamma = tensor(Shape(4), gammaArr), + initBeta = tensor(Shape(4), betaArr), + ) + + val output = gn.forward(input, ctx) + + val expected = referenceGroupNorm(data, n = 1, c = 4, spatial = 4, groups = 2, gamma = gammaArr, beta = betaArr) + assertFlatClose(expected, output, Shape(1, 4, 2, 2)) + } + + @Test + fun groupNorm_defaultAffine_isIdentityScaleShift() { + // Default placeholder params are gamma=1, beta=0 (the `groupNorm {}` DSL path): + // output must equal the non-affine normalization. + val data = floatArrayOf(2f, 4f, 6f, 8f, 1f, 3f, 5f, 7f) + val withDefaultAffine = GroupNormalization( + numGroups = 2, numChannels = 4, name = "gn_default", dtype = FP32::class, + ) + val withoutAffine = GroupNormalization( + numGroups = 2, numChannels = 4, affine = false, name = "gn_plain", + ) + + val a = withDefaultAffine.forward(tensor(Shape(2, 4), data), ctx) + val b = withoutAffine.forward(tensor(Shape(2, 4), data), ctx) + + for (row in 0..1) { + for (col in 0..3) { + assertAlmostEquals(b.data[row, col] as Float, a.data[row, col] as Float, 1e-6f) + } + } + } + + @Test + fun groupNorm_singleGroup_normalizesWholeSample() { + // G=1 normalizes over all C*spatial elements of each sample (LayerNorm-like). + val data = floatArrayOf( + 1f, 2f, 3f, 4f, // sample 0 + 10f, 20f, 30f, 40f, // sample 1 + ) + val input = tensor(Shape(2, 4), data) + val gn = GroupNormalization(numGroups = 1, numChannels = 4, affine = false, name = "gn1") + + val output = gn.forward(input, ctx) + + val expected = referenceGroupNorm(data, n = 2, c = 4, spatial = 1, groups = 1) + assertFlatClose(expected, output, Shape(2, 4)) + // Each sample has ~zero mean after normalization + for (row in 0..1) { + var sum = 0f + for (col in 0..3) sum += output.data[row, col] as Float + assertAlmostEquals(0f, sum / 4f, 1e-5f, "Sample $row mean should be ~0") + } + } + + @Test + fun groupNorm_groupsEqualChannels_isInstanceNorm() { + // G=C: each channel normalized over its own spatial elements only. + val data = floatArrayOf( + 1f, 3f, // channel 0: mean 2, var 1 + 10f, 30f, // channel 1: mean 20, var 100 + ) + val input = tensor(Shape(1, 2, 2), data) + val gn = GroupNormalization(numGroups = 2, numChannels = 2, affine = false, name = "gnC") + + val output = gn.forward(input, ctx) + + val expected = referenceGroupNorm(data, n = 1, c = 2, spatial = 2, groups = 2) + assertFlatClose(expected, output, Shape(1, 2, 2)) + // Hand math: (1-2)/sqrt(1+1e-5) ~= -0.999995, (10-20)/sqrt(100+1e-5) ~= -1.0 + assertAlmostEquals(-0.999995f, output.data[0, 0, 0] as Float) + assertAlmostEquals(-1f, output.data[0, 1, 0] as Float) + } + + @Test + fun groupNorm_constantInput_normalizesToZero() { + // Zero variance: (x - mean) / sqrt(0 + eps) = 0 (eps prevents NaN/Inf). + val input = tensor(Shape(1, 4), floatArrayOf(5f, 5f, 5f, 5f)) + val gn = GroupNormalization(numGroups = 2, numChannels = 4, affine = false, name = "gn_const") + + val output = gn.forward(input, ctx) + + for (i in 0..3) { + assertAlmostEquals(0f, output.data[0, i] as Float, 1e-2f, "Constant input should normalize to ~0") + } + } + + @Test + fun groupNorm_customEps_entersDenominator() { + // Large eps gives a visibly different result: (x - m) / sqrt(var + 1.0) + val data = floatArrayOf(1f, 2f, 3f, 4f) + val input = tensor(Shape(1, 4), data) + val gn = GroupNormalization(numGroups = 2, numChannels = 4, eps = 1.0, affine = false, name = "gn_eps") + + val output = gn.forward(input, ctx) + + val expected = referenceGroupNorm(data, n = 1, c = 4, spatial = 1, groups = 2, eps = 1.0) + assertFlatClose(expected, output, Shape(1, 4)) + // var 0.25, denom sqrt(1.25): clearly distinct from the eps=1e-5 result + assertAlmostEquals((-0.5 / sqrt(1.25)).toFloat(), output.data[0, 0] as Float) + } + + @Test + fun groupNorm_wrongChannelCount_throws() { + val input = tensor(Shape(1, 6), floatArrayOf(1f, 2f, 3f, 4f, 5f, 6f)) + val gn = GroupNormalization(numGroups = 2, numChannels = 4, affine = false, name = "gn_bad") + + assertFailsWith { gn.forward(input, ctx) } + } + + @Test + fun groupNorm_indivisibleGroups_throwsAtConstruction() { + assertFailsWith { + GroupNormalization(numGroups = 3, numChannels = 4, name = "gn_indiv") + } + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/normalization/GroupNormalization.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/normalization/GroupNormalization.kt index 231d395af..19fd8646a 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/normalization/GroupNormalization.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/normalization/GroupNormalization.kt @@ -4,11 +4,13 @@ import sk.ainet.context.ExecutionContext import sk.ainet.lang.nn.Module import sk.ainet.lang.nn.topology.ModuleParameter import sk.ainet.lang.nn.topology.ModuleParameters -import sk.ainet.lang.ops.NotImplemented import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.* +import sk.ainet.lang.types.BF16 import sk.ainet.lang.types.DType +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 import kotlin.reflect.KClass /** @@ -23,7 +25,6 @@ import kotlin.reflect.KClass * @param initGamma Initial gamma (scale) parameter * @param initBeta Initial beta (shift) parameter */ -@NotImplemented(reason = "forward helpers (group reshape, mean, variance) are TODO(); eps and gamma/beta are not applied", issue = "#1379") public class GroupNormalization( private val numGroups: Int, private val numChannels: Int, @@ -86,55 +87,64 @@ public class GroupNormalization( override fun forward(input: Tensor, ctx: ExecutionContext): Tensor = sk.ainet.lang.nn.hooks.withForwardHooks(ctx, this, input) { - // Reshape input to separate groups: (N, C, H, W) -> (N, G, C//G, H, W) - val groupedInput = reshapeForGroups(input) - - // Calculate group statistics - val groupMean = calculateGroupMean(groupedInput) - val groupVar = calculateGroupVariance(groupedInput, groupMean) - - // Normalize within groups - val normalized = normalizeGroups(groupedInput, groupMean, groupVar) - - // Reshape back to original format - val output = reshapeFromGroups(normalized, input.shape) - + val inputShape = input.shape + require(inputShape.rank >= 2) { + "GroupNormalization expects input of shape (N, C, ...), got rank ${inputShape.rank}" + } + require(inputShape.dimensions[1] == numChannels) { + "Input has ${inputShape.dimensions[1]} channels, expected $numChannels" + } + + // Compute the numerically-sensitive normalization (mean / variance / + // sqrt / divide) in f32 regardless of the model dtype, mirroring + // LayerNormalization: a low-precision variance loses enough precision + // that `sqrt(var + eps)` can overflow to NaN. The affine (gamma/beta) + // stays in the model dtype. No-op when the model already runs f32. + val normalized: Tensor = if (input.dtype != FP32::class) { + val xF32 = input.ops.convert(input, FP32) + val normF32 = normalizeGroups(xF32) + val modelDtype: DType = when (input.dtype) { + BF16::class -> BF16 + FP16::class -> FP16 + else -> FP32 + } + @Suppress("UNCHECKED_CAST") + (input.ops.convert(normF32, modelDtype) as Tensor) + } else { + normalizeGroups(input) + } + if (affine) { - val gamma = params[0].value // weight parameter - val beta = params[1].value // bias parameter - output * gamma + beta + // gamma/beta are per-channel (C); reshape to (1, C, 1, ...) so they + // broadcast over the batch and spatial dimensions. The reshape must run + // on the input's ops: the default placeholder params carry VoidTensorOps, + // whose reshape is shape-only and would zero the values. + val broadcastShape = Shape(IntArray(inputShape.rank) { idx -> if (idx == 1) numChannels else 1 }) + val gamma = input.ops.reshape(params[0].value, broadcastShape) + val beta = input.ops.reshape(params[1].value, broadcastShape) + normalized * gamma + beta } else { - output + normalized } } - private fun reshapeForGroups(input: Tensor): Tensor { - // Reshape from (N, C, H, W) to (N, G, C//G, H, W) - TODO("Implement group reshaping") - } - - private fun calculateGroupMean(groupedInput: Tensor): Tensor { - // Calculate mean across spatial and channel dimensions within each group - TODO("Implement group mean calculation") - } + /** + * Normalizes with `torch.nn.GroupNorm` semantics: (N, C, *) is reshaped to + * (N, G, C/G * prod(*)), each group is normalized with its own biased + * mean/variance over the last dimension, and the result is reshaped back. + */ + private fun normalizeGroups(input: Tensor): Tensor { + val originalShape = input.shape + var rest = 1 + for (i in 2 until originalShape.rank) rest *= originalShape.dimensions[i] + val grouped = input.reshape(Shape(originalShape.dimensions[0], numGroups, channelsPerGroup * rest)) - private fun calculateGroupVariance(groupedInput: Tensor, groupMean: Tensor): Tensor { - // Calculate variance across spatial and channel dimensions within each group - TODO("Implement group variance calculation") - } - - private fun normalizeGroups( - groupedInput: Tensor, - groupMean: Tensor, - groupVar: Tensor - ): Tensor { - // normalized = (input - mean) / sqrt(variance) - // Note: eps handling would need proper scalar-tensor operations - return (groupedInput - groupMean) / groupVar.sqrt() - } + val mean = grouped.mean(dim = 2).unsqueeze(2) // (N, G, 1) + val centered = grouped - mean + // Biased variance (mean of squared deviations), matching torch.nn.GroupNorm. + val variance = (centered * centered).mean(dim = 2).unsqueeze(2) + val normalized = centered / (variance + eps).sqrt() - private fun reshapeFromGroups(normalized: Tensor, originalShape: Shape): Tensor { - // Reshape from (N, G, C//G, H, W) back to (N, C, H, W) - TODO("Implement reshape from groups") + return normalized.reshape(originalShape) } } \ No newline at end of file