Repository navigation
feat(nn): implement GroupNormalization forward pass (#1381) - #1463
Merged
Merged
Conversation
Implements the four TODO() forward helpers 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, eps enters the denominator, and the result is reshaped back. - Numerically-sensitive normalization runs in f32 regardless of model dtype, mirroring LayerNormalization. - Per-channel gamma/beta are reshaped to (1, C, 1, ...) through the input's ops for broadcasting; the default placeholder params carry VoidTensorOps, whose shape-only reshape would zero the values. - Input rank and channel count are validated with require(). - @NotImplemented removed; the layer leaves the coverage-matrix gap list. Tests in skainet-backend-cpu commonTest compare against a double- precision reference implementation of torch.nn.GroupNorm covering 2D/3D/4D inputs, per-channel affine, single group, groups == channels (instance norm), constant input, custom eps and validation failures. Part of #1379.
|
📖 Documentation Preview The documentation has been built successfully for this PR. Generated Files:
Artifacts:
This comment will be updated automatically when the PR is updated. |
GroupNormalization moves from the NN-layers gap list to done (31/32). Regenerated with the HLO and backward coverage dumps per the documentation block in build.gradle.kts.
|
📖 Documentation Preview The documentation has been built successfully for this PR. Generated Files:
Artifacts:
This comment will be updated automatically when the PR is updated. |
aharakal
approved these changes
Oct 31, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements Lane 2 (kotlin-core) sub-issue #1381 of the GroupNormalization feature #1379.
GroupNormalizationwas reachable from the publicgroupNorm {}DSL builder but crashed on the first forward pass: its four forward helpers wereTODO(),epswas never applied, and gamma/beta were never broadcast. This PR implements the forward pass withtorch.nn.GroupNormsemantics using ops that already exist.Part of #1379.
Closes #1381.
Changes
Forward pass (
sk.ainet.lang.nn.normalization.GroupNormalization):(N, C, *)is reshaped to(N, G, C/G * prod(*)); each group is normalized with its own biased mean/variance (matching PyTorch),epsenters the denominator as(x - mean) / sqrt(var + eps), and the result is reshaped back.LayerNormalizationin the same package.(1, C, 1, ...)for broadcasting over batch and spatial dimensions. The reshape runs on the input's ops: the default placeholder params carryVoidTensorOps, whose shape-onlyreshapereturns zeros and would silently wipe the affine parameters (the same latent hazard exists inBatchNormalization, which reshapes stats through data-bearing tensors only).require().@NotImplementedremoved — the layer leaves the ❌ list in the coverage matrix.Tests (
skainet-backend-cpu/src/commonTest/.../GroupNormalizationTest.kt, 10 tests):torch.nn.GroupNorm(per-group biased stats + per-channel affine), within 1e-5.(N, C), 3D and 4D inputs; per-channel gamma/beta broadcast (distinct values per channel, the bug-prone axis); default placeholder affine (= thegroupNorm {}DSL path) equals non-affine output;G = 1(layer-norm-like) andG = C(instance-norm-like) extremes; constant input stays finite via eps; custom eps enters the denominator; wrong channel count and indivisible groups throwIllegalArgumentException.Checks run locally
:skainet-lang:skainet-lang-core:jvmTestand:skainet-backends:skainet-backend-cpu:jvmTest— green:skainet-lang:skainet-lang-core:jvmApiCheck— no public API shape changecheckUnmarkedStubs— green (no allowlist entry needed; all TODOs removed)Note
Lane 1 (#1380, PyTorch reference values) has no posted fixtures yet; the tests encode the reference formula directly in double precision, which is bit-equivalent to
torch.nn.GroupNormfor these fixture sizes. Happy to swap in posted fixtures when #1380 lands. Docs (#1382) and review (#1383) remain as separate lanes.