fix(dnn): BatchNorm running-var bias, pooling channel_axis, LayerNorm error

[chaoming0625]

Fresh review of brainpy/dnn. Three Medium fixes:

• BatchNorm stored the biased batch variance into running_var; now applies Bessel's N/(N-1) correction for the running buffer (PyTorch-consistent), keeping biased variance for in-batch normalization.
• Pooling rejected the leftmost negative channel_axis (== -x_dim) due to an abs() bound; widened to -x_dim <= axis < x_dim (Pool, _MaxPoolNd, AdaptivePool).
• LayerNorm wrong-shape error did ', '.join(<ints>) -> TypeError masking the intended ValueError.

Prior-audit GroupNorm/InstanceNorm + affine-norm findings verified already-fixed. Regression tests added (589 passed in-scope). Findings: docs/issues-found-20260619-dnn.md.

Summary by Sourcery

Fix dnn normalization and pooling edge-case bugs and document the audit findings for these modules.

Bug Fixes:

• Ensure LayerNorm raises a clear ValueError with a correct expected-shape message when the input trailing dimensions mismatch normalized_shape.
• Update BatchNorm to use an unbiased (Bessel-corrected) variance estimate for the running_var buffer while preserving biased variance for in-batch normalization.
• Allow pooling layers to accept the leftmost negative channel_axis (== -x_dim) by using a symmetric axis bounds check instead of abs().

Documentation:

• Add an issues-found report documenting the dnn module audit, identified problems, and their resolution status.

Tests:

• Add regression tests covering LayerNorm shape-mismatch errors, BatchNorm running_var unbiasedness and batch normalization behavior, and pooling with leftmost negative channel_axis.

[sourcery-ai]

Reviewer's Guide

Adjusts BatchNorm to store an unbiased (Bessel-corrected) running variance while keeping biased in-batch normalization, fixes pooling channel_axis validation to accept the leftmost negative axis, and corrects LayerNorm’s shape-mismatch error message to raise the intended ValueError; adds focused regression tests and an audit doc describing the issues and their status.

Flow diagram for BatchNorm.update unbiased running_var computation

flowchart TD
  X["Input x"] --> A["Compute mean over axis"]
  A --> B["Compute mean_of_square over axis"]
  B --> C["var = max(0, mean_of_square - mean^2)<br/>(biased, divisor N)"]
  C --> D["num_reduced = product of reduced x.shape[axis]"]
  D --> E{"num_reduced > 1?"}
  E -- Yes --> F["unbiased_var = var * (num_reduced / (num_reduced - 1))"]
  E -- No --> G["unbiased_var = var"]
  F --> H["running_var = momentum * running_var + (1 - momentum) * unbiased_var"]
  G --> H
  C --> I["Use biased var for in-batch normalization"]

Loading

Flow diagram for pooling channel_axis validation and normalization

flowchart TD
  X["Given x_dim and channel_axis"] --> A{"channel_axis is not None and != 0?"}
  A -- No --> Z["Use default channel_axis"]
  A -- Yes --> B{"-x_dim <= channel_axis < x_dim?"}
  B -- No --> C["Raise ValueError: Invalid channel axis"]
  B -- Yes --> D{"channel_axis < 0?"}
  D -- Yes --> E["channel_axis = x_dim + channel_axis<br/>(convert negative to positive index)"]
  D -- No --> F["Keep channel_axis as is"]
  E --> G["Proceed with pooling using resolved channel_axis"]
  F --> G

Loading

File-Level Changes

Change	Details	Files
BatchNorm running variance now uses an unbiased (Bessel-corrected) estimate while preserving biased normalization for the current batch.	Compute batch variance as before for normalization, then derive an unbiased variance by scaling with N/(N-1) over all reduced axes, guarding the N==1 case. Update running_var using the unbiased variance in the EMA, while running_mean logic remains unchanged. Add regression tests verifying that running_var matches the unbiased update and that the normalized batch still has mean≈0 and (biased) var≈1.	brainpy/dnn/normalization.py brainpy/dnn/normalization_test.py
Pooling layers now accept the leftmost negative channel_axis (== -x_dim / -x.ndim) instead of rejecting it due to an abs()-based bound check.	Change channel_axis validation in Pool._infer_shape, _MaxPoolNd._infer_shape, and AdaptivePool.update from an abs()-based bound to a numpy-style range check -x_dim <= channel_axis < x_dim (or -x.ndim <= channel_axis < x.ndim). Retain existing behavior of translating negative channel_axis values to their positive equivalents after validation. Add regression tests that exercise MaxPool2d, AdaptiveAvgPool2d, and Pool-family MaxPool with channel_axis=-3 on (C,H,W) inputs and assert expected output shapes and equivalence to the corresponding positive axis.	brainpy/dnn/pooling.py brainpy/dnn/pooling_layers_test.py
LayerNorm’s shape-mismatch guard now constructs its error message correctly and reliably raises ValueError instead of being masked by a TypeError.	Fix the expected-shape formatting in LayerNorm.update to join normalized_shape via map(str, ...) so the ValueError message can be constructed for non-string items. Keep the mismatch condition itself unchanged, only adjusting message construction. Add a regression test that feeds an input whose trailing dimension does not match normalized_shape and asserts that ValueError is raised.	brainpy/dnn/normalization.py brainpy/dnn/normalization_test.py
Document the broader dnn audit findings and cross-checks in a new markdown report.	Add docs/issues-found-20260619-dnn.md summarizing the expert review scope, environment, severity levels, concrete issues (P12-M1/M2/M3, P12-L1), their status, and cross-checks against a prior audit. Record several components as already-correct with no code changes (e.g., GroupNorm/InstanceNorm, Dropout, Conv/ConvTranspose, activation wrappers, interoperation_flax).	docs/issues-found-20260619-dnn.md

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with @sourcery-ai issue to create an issue from it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time. You can also comment
  @sourcery-ai title on the pull request to (re-)generate the title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment @sourcery-ai summary on the pull request to
  (re-)generate the summary at any time.
• Generate reviewer's guide: Comment @sourcery-ai guide on the pull
  request to (re-)generate the reviewer's guide at any time.
• Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
• Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[sourcery-ai]

Hey - I've reviewed your changes and they look great!

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}