Skip to content

Fix BatchNorm docstring to cover 4D input and test the NHWC path#3782

Merged
zcbenz merged 1 commit into
ml-explore:mainfrom
Pablosinyores:fix/batchnorm-doc-4d
Jul 1, 2026
Merged

Fix BatchNorm docstring to cover 4D input and test the NHWC path#3782
zcbenz merged 1 commit into
ml-explore:mainfrom
Pablosinyores:fix/batchnorm-doc-4d

Conversation

@Pablosinyores

Copy link
Copy Markdown
Contributor

Summary

nn.BatchNorm's docstring summary says it applies over a "2D or 3D input", but the layer also supports 4D NHWC tensors:

  • BatchNorm.__call__ guards x.ndim < 2 or x.ndim > 4 and its own error message says "Expected input tensor to have 2, 3 or 4 dimensions".
  • The docstring body already documents the 4D case: "For four-dimensional arrays, the shape is NHWC...".

So the summary line contradicts both the code and the rest of its own docstring. This corrects the summary to "2D, 3D or 4D input".

Test

test_batch_norm previously covered only 2D and 3D inputs (plus a 5D ValueError check); the 4D NHWC path was never exercised. Added a 4D case asserting the output is normalized per channel across N, H and W (per-channel mean ≈ 0, var ≈ 1).

No behavior change — documentation + test coverage only.

@zcbenz zcbenz left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

The BatchNorm summary line said it applies over a 2D or 3D input, but the
layer also accepts 4D NHWC tensors: the __call__ guard allows 2-4 dims (and
its error message says "2, 3 or 4 dimensions"), and the docstring body
already describes the four-dimensional NHWC shape. Correct the summary to
match, and add the missing 4D coverage to test_batch_norm, asserting that an
NHWC input is normalized per channel across N, H and W.
@zcbenz zcbenz force-pushed the fix/batchnorm-doc-4d branch from 05881b5 to a105e98 Compare June 30, 2026 23:39
@zcbenz zcbenz merged commit e9463bb into ml-explore:main Jul 1, 2026
21 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants