[sylvesterkaczmarek]

Compute GroupNorm's reshape dimensions from the current input rather than caching the first input's spatial dimensions. Reusing one module within a forward pass on (1, 2, 3, 8) followed by (1, 3, 2, 8) currently returns the second result with shape (1, 2, 3, 8). Changing the spatial volume can instead change the inferred batch size or raise a reshape error.

The fix keeps the existing rank check, reduction-axis mapping, parameter names, and per-channel parameter shapes. Both channel layouts continue to share one set of scale/offset parameters. This restores the documented same-shape output contract for repeated use at different spatial resolutions.

Newly reproduced on current main e4dbc6e. Searches of issues and open PRs found no matching fix or ownership claim. This is independent of the pending attention and BatchNorm changes.

Validation

macOS arm64 CPU, Python 3.12.11:

• New regression module: 24 failures and 8 batch-size controls before the fix; all 32 cases pass afterward. Coverage includes both channel layouts, changed spatial layout and volume, parameter sharing, affine/non-affine normalization, eager/JIT execution, and input gradients against an independent channel-group reference.
• New and existing GroupNorm suites: 60 passed.
• Main non-integration suite: python -m pytest -n 2 -q haiku --ignore=haiku/_src/integration --tb=short --maxfail=5 completed with 1,591 passed, 34 skipped, 18 passing subtests, and one failure in the unrelated ReshapeTest.test_reshape_convert export test.
• That export failure was reproduced on a pristine upstream worktree: the installed TensorFlow cannot deserialize the newer StableHLO artifact produced by JAX. No test was removed or weakened to hide it.
• Source distribution and wheel builds pass. Installed the wheel outside the checkout, verified its import path, and reran all 60 GroupNorm tests successfully. The 32 regressions also pass with JAX's x64 mode enabled.
• Ruff correctness lint, new-test Pyink formatting, syntax compilation, and git diff --check pass. Changed-file Pylint errors-only exits successfully; its existing configuration emits obsolete-option diagnostics.

GPU/TPU execution, the integration directory, documentation builds, and other operating systems were not tested locally. No existing assertion, tolerance, dependency, or CI configuration was changed.

Validation update on 7 October 2026

Current validated commit d8c8aed42e9b.

Refreshed onto checked current main. Existing component tests and the new regression file pass in the normal default-precision configuration using real JAX CPU execution. 60 passed in 20.30s Pylint error checks and git diff --check pass. Only the new test file was formatted; no whole-source rewrite of existing main formatting was added. The complete repository and GPU/TPU matrix were not run locally.

The complete upstream workflow result is not claimed from local checks. Fresh fork workflows may require maintainer approval.

[sylvesterkaczmarek]

Updated validation is recorded in the PR description for d8c8aed42e9b. The branch is refreshed onto the checked live base. Could a maintainer review this head and approve any waiting workflows?

[sylvesterkaczmarek]

Following up on #907, #906, #905, #914, #913, #911. Could a maintainer take a look when available? Please let me know if any need changes or a different scope.